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

## Abstract

Abstract Because it is difficult to precisely specify complex objectives, reinforcement learning policies are often optimized using flawed proxy rewards that seem to capture the true objective. However, optimizing proxy rewards frequently leads to reward hacking : the optimized reward function ceases to be a good proxy, and the resulting policy performs poorly with respect to the unspecified true reward. Principled solutions to reward hacking have been impeded by the lack of a good definition for the problem. To address this, we introduce a definition of reward hacking based on the correlation between proxy and true rewards for states and actions seen by a “base policy” that breaks down under optimization. We show that this definition captures reward hacking behavior across several realistic settings, including in reinforcement learning from human feedback (RLHF). We then show theoretically that regularization to the base policy can effectively prevent reward hacking. While current RLHF approaches apply a KL penalty between the action distributions of policies, our theory suggests that it is more effective to regularize using the χ 2 \chi^{2} divergence between the policies’ occupancy measures. We intuitively show why this type of regularization is superior and demonstrate that it better mitigates reward hacking in practice across four realistic domains, including RLHF for LLMs. Our code is available at https://github.com/cassidylaidlaw/orpo .

## 1 Introduction

A major challenge for the designers of goal-oriented AI systems is specifying a reward function that robustly captures their goals and values. Manually designing reward functions is difficult due to the ambiguities and complexity underlying real-world scenarios (Ibarz et al., ). An alternative is to learn reward functions from human data (Sadigh et al., ; Jeon et al., ), but these often fail to generalize outside the distribution of behavior seen during training (McKinney et al., ; Tien et al., ).
Thus, a learned or hand-specified reward function is often just a *proxy* for the true reward underlying the system designer’s intent. Misalignment between the two objectives can lead to *reward hacking*: a learned policy performs well according to the proxy reward function, but not according to the true reward function (Russell et al., ; Amodei et al., ; Pan et al., ; Skalse et al., ). A reward hacking policy’s behavior is often undesirable and can be especially catastrophic when deployed in safety-critical scenarios, such as autonomous driving (Krakovna et al., ; Turner et al., ; Knox et al., ). Unfortunately, reward hacking is a common phenomenon (Krakovna, ) and has harmful effects in the real world (Lum & Isaac, ; Corbett-Davies et al., ; Obermeyer et al., ; Milli et al., ; Franchi et al., ; Kleinberg et al., ).

The ideal solution to prevent reward hacking would be to perfectly align the specified proxy and unknown true reward; however, in many domains, this is impossible to achieve. For example, imagine trying to design a reward function for a self-driving car. It would have to capture the speed at which the car reached the destination, comfort of the passenger, all applicable laws, and other factors, many of which are difficult to robustly measure; furthermore, these factors would have to be carefully weighted against each other (Knox et al., ).
In practice, reward hacking can occur even with significant reward engineering efforts. For example, hand-designed reward functions for recommender systems have led to adverse outcomes in terms of user experience (Stray et al., ). Reinforcement learning from human feedback (RLHF), where a learned reward function is used to align a large-language model (LLM) based assistant to human preferences, provides another example. Without regularization, RLHF leads to the overoptimization of the learned reward function, a case of reward hacking (Gao et al., ).

Since proxy reward functions for complex tasks are nearly always misspecified in practice, what can be done to avoid reward hacking? There is a lack of principled solutions for preventing reward hacking, stemming from a more fundamental problem: *defining* reward hacking in a formal sense that captures realistic cases.
In particular, it is difficult to characterize what makes a proxy reward function “good”. For example, Skalse et al. () define a notion of “unhackability” for proxy reward functions but find that it can only hold if the proxy is a multiple of the true reward, which is obviously unrealistic.

Figure: Figure 1: We present a new characterization of reward hacking and a method for preventing it. We define a proxy reward function as one that correlates with an unknown true reward function for state-action pairs sampled from some base policy. However, optimizing the proxy alone can lead to a breakdown in the correlation and worse true reward than the base policy. We show theoretically and empirically that optimizing the proxy with $\chi^{2}$ occupancy measure regularization to the base policy can allow outperforming the base policy under the unknown true reward.

We argue that proxies are chosen because they seem to capture the true objective under some base distribution of behavior. To formalize this, we define a proxy reward as one that *correlates* with the true reward function under the distribution of states and actions visited by a *base policy*. For example, if we use a human-like policy as the base policy for self-driving cars, then a proxy reward would have to correlate with the true reward for states and actions seen during typical human driving. In the case of RLHF, where the reward function is learned from human preference feedback over different model responses, a natural base policy is the supervised fine-tuning (SFT) model used to generate those responses. If the reward learning process is effective, then the learned reward should correlate well with the true human-judged reward on the distribution of responses it was trained on, making it a good proxy.

We thus formally define reward hacking as occurring when optimizing a proxy reward induces a decrease in true reward compared to the base policy (Figure 1). We show that this definition and our definition of a proxy reward function capture a number of intuitive cases of reward hacking in realistic environments, including traffic control, blood glucose regulation, and RLHF (Figure 2).

Furthermore, we show that our definition allows us to derive a method for avoiding reward hacking during policy optimization by regularizing the optimized policy to be similar to the base policy. Specifically, we find that optimizing the proxy reward minus a regularization term provides a *provable* lower bound on improvement in the true reward. The amount of regularization needed depends on how strong the correlation is between the proxy and true rewards; as the correlation becomes stronger, it is possible to use less regularization.

Regularizing policy optimization to a base policy is already used in practice during RLHF via a KL divergence penalty (Stiennon et al., ; Bai et al., ), and Proposition 5.1 provides some theoretical justification for why this works. However, while the regularization in RLHF penalizes the KL divergence between the optimized and base policies’ *action distributions*, our result suggests that it is better to penalize the chi-squared ($\chi^{2}$) divergence between the policies’ *occupancy measures*. A policy’s occupancy measure (OM) is the distribution of states or state-action pairs seen by the policy when it interacts with its environment. Unlike action distribution (AD)-based metrics, occupancy measures take into account the states that the agent reaches, not just the actions it takes. While algorithms based on occupancy measures have been widely used for imitation learning (Ho & Ermon, ), offline RL (Lee et al., ), and efficient exploration (Hazan et al., ), using OM divergences to prevent reward hacking remains unexplored. We compare OM vs. AD regularization and $\chi^{2}$ vs. KL divergence, presenting intuitive reasons why $\chi^{2}$ OM divergence may be a better regularization target in Figures 3 and 4.

Based on our theoretical results, we empirically investigate using the benefits of using occupancy measure regularization to prevent reward hacking.
We implement OM-based regularization in practice using a discriminator network that approximates the OM divergence between policies.
We then optimize policies with misaligned proxy reward functions in multiple reward hacking benchmark environments (Pan et al., ) using OM and AD regularization.
The results of our experiments demonstrate that training with occupancy measure regularization leads to better performance under the unseen true reward function in all of the environments, validating our theoretical results.
In contrast, we find that it is difficult to tune AD regularization in some environments to both prevent reward hacking and allow meaningful improvement over the base policy.
Furthermore, we compare regularizing with $\chi^{2}$ divergence, which our theoretical results suggest using, to KL divergence, the prior standard in settings like RLHF. We find that regularization with chi-squared divergence leads to more stable results across regularization coefficients and often achieves higher true reward.

Our main contributions can be summarized as follows:

- 1.
We provide a new formal definition of reward hacking and show that it captures a number of realistic case studies.
- 2.
Using our definition, we establish that optimizing a proxy reward with $\chi^{2}$ occupancy measure regularization leads to a provable improvement in the unknown true reward function.
- 3.
We implement $\chi^{2}$ OM regularization in practice and demonstrate that it is superior to the current standard for preventing reward hacking via regularization.

## 2 Related work

While there have been separate lines of work investigating reward hacking and exploring the use of occupancy measure divergences for other applications, to the best of our knowledge, we are the first to specifically study applying occupancy measure regularization to the problem of reward hacking.

Reward hacking.
Some prior works establish theoretical models of reward hacking as a special case of Goodhart’s Law (Goodhart, ; Leike et al., ; Manheim & Garrabrant, ; Krakovna, ; Skalse et al., ; Ngo et al., ).
Krakovna () provide a list of many examples of reward hacking. Pan et al. () categorize types of reward misspecification and relate optimization power to reward hacking.

Safe reinforcement learning.
Regularizing policies to be similar to an offline policy based on their action distribution KL divergence was proposed by Stiennon et al. () and has since been widely employed in the context of optimizing LLMs using RLHF (Ouyang et al., ; Bai et al., ; Glaese et al., ). KL regularization for RLHF has been further studied by Vieillard et al. (), Gao et al. (), and Korbak et al. ().
Some alternative approaches to avoid reward hacking include quantilizers (Taylor, ), “mild” optimization (Taylor et al., ), and impact regularization (Turner et al., ).
While constrained RL can prevent the misbehavior of agents that optimize flawed reward functions (Dalal et al., ; Chow et al., ; Zhang et al., ; Roy et al., ), it simply shifts the difficulty of designing a reward function to specifying a set of constraints and weights.
Robust RL usually considers a misspecified transition model, but some work has explored misspecified reward functions (Derman et al., ; Gadot et al., ).
Other proposals to address the reward specification problem attempt to infer the true reward function based on the given proxy reward function, environment context, and/or feedback from humans (Hadfield-Menell et al., ; Reddy et al., ; Lee et al., ).

Applications of occupancy measure regularization.
Occupancy measure regularization and optimization have been used for a variety of purposes. GAIL (Ho & Ermon, ) is an algorithm for robust imitation learning that aims to match the imitator’s occupancy measure to that of the demonstrator. Kang et al. () combines GAIL with a reward function to efficiently explore using human data. Another line of work aims to find a policy with the highest-entropy occupancy measure for the purpose of exploring the state space (Hazan et al., ; Lee et al., ; Nedergaard & Cook, ). Various types of distributional regularization are used in model-based RL since learned models may not generalize out-of-distribution (Yang et al., ).

Offline reinforcement learning.  Prior work in offline RL may appear to be particularly related to our work. Many offline RL algorithms use occupancy measure or action distribution-based regularization to ensure that the learned policy remains within the training data distribution (Lee et al., ; Mandal et al., ; He, ; Cheng et al., ; Rashidinejad et al., ; Xie et al., ). However, the settings are fundamentally different: while offline RL is limited by a lack of coverage in the training data, the difficulty in our setting is that the reward function is misspecified. While it might be possible to avoid reward hacking by using offline RL algorithms, we leave this to future work and focus on regularization-based approaches in the online RL setting.

Figure: Figure 2: We show that reward hacking in four realistic environments and an illustrative gridworld is captured by our definition. The top row shows the distribution of proxy and true reward values for state-action pairs sampled from a natural base policy for each environment; in all environments, the proxy and true rewards are correlated. However, if the proxy is optimized via reinforcement learning, the correlation breaks down and the true reward is lower than the base policy, which we define as reward hacking (middle row). Our theoretical results show that RL with occupancy measure regularization to the base policy can prevent reward hacking and enable an increase in true reward (bottom row).
Refer to caption: figures/traffic_env.png

## 3 Preliminaries

To study reward hacking, we consider the setting of an infinite-horizon Markov decision process (MDP). An agent takes actions $a\in\mathcal{A}$ to transition between states $s\in\mathcal{S}$ over a series of timesteps $t=0,1,2,\ldots$. We assume that $\mathcal{S}$ and $\mathcal{A}$ are finite for simplicity but our results can easily generalize to infinite state or action spaces. The first state $s_{0}$ is sampled from an initial distribution $\mu_{0}(s)$, and when an agent takes action ${a_{t}}$ in ${s_{t}}$ at time ${t}$, the next state $s_{t+1}$ is reached at timestep $t+1$ with transition probability $p(s_{t+1}\mid s_{t},a_{t})$. The agent aims to optimize a reward function $R:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$, and rewards are accumulated over time with discount factor $\gamma\in[0,1)$. A policy $\pi$ maps each state $s$ to a distribution over actions to take at that state $\pi(a\mid s)$. We define the (normalized) *return* of a policy $\pi$ under a reward function $R$ as

$$ $\textstyle J(\pi,R)=(1-\gamma)\,\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{t})\right]$ (1) $$

where $\mathbb{E}_{\pi}$ refers to the expectation under the distribution of states and actions induced by running the policy $\pi$ in the environment.

We define the *state-action occupancy measure* $\mu_{\pi}$ of a policy $\pi$ as the expected discounted number of times the agent will be in a particular state and take a specific action:

$$ $\textstyle\mu_{\pi}(s,a)=(1-\gamma)\,\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}\mathds{1}\{s_{t}=s\wedge a_{t}=a\}\right].$ (2) $$

Intuitively, the occupancy measure represents the distribution of states and actions visited by the policy over time. Furthermore, combining (1) and (2), it is easy to show that the return of a policy can be re-written as the expected reward for states and actions sampled from the occupancy measure:

$$ $\textstyle J(\pi,R)=\sum_{s,a\in\mathcal{S}\times\mathcal{A}}\mu_{\pi}(s,a)R(s,a)=\mathbb{E}_{\mu_{\pi}}[R(s,a)].$ (3) $$

The standard approach to solving an MDP is to find a policy $\pi$ that maximizes its return $J(\pi,R)$.
However, as we discussed in the introduction, an AI system designer might optimize $\pi$ using a learned or hand-specified *proxy* reward function $\tilde{R}$ which is misaligned with the *true* reward function $R$. In order to better understand and mitigate reward hacking, it would be helpful to have a good definition of the problem. Intuitively, reward hacking is when optimizing the proxy reward $J(\pi,\tilde{R})$ of a policy $\pi$ ultimately leads to low true reward $J(\pi,R)$. However, this intuition is more difficult to formalize than it might seem.

Desiderata for a definition of reward hacking.
To understand why defining reward hacking is difficult, consider the case study of RLHF. In this case, optimizing a learned reward function over LLM outputs eventually leads to the LLM producing nonsensical responses.
This clearly satisfies our informal definition of “optimizing the proxy makes the true reward go down.” However, what if we were to optimize $\tilde{R}(s,a)=-R(s,a)$? In this case, optimizing $\tilde{R}$ also makes the true reward go down, but arguably this is not reward hacking because $\tilde{R}$ was not a good “proxy” in the first place.
Thus, a good definition of reward hacking should distinguish between “reasonable” proxies and reward functions that are clearly unrelated (or opposite) to the true reward function.

In the RLHF example, optimizing until the LLM produces nonsensical outputs seems like bad enough behavior that we would say the proxy has been “hacked.”
However, what if optimizing the proxy produced mediocre but not obviously bad outputs? Even if optimizing the proxy does not lead to near-optimal true reward, we would not always say that reward hacking has occurred. Thus, a good definition of reward hacking should also choose a threshold for true reward: below the threshold, reward hacking is occurring, and above the threshold it is not.

Prior definitions of reward hacking.
Because of these difficulties, prior work has struggled to define reward hacking. With regard to our first desideratum, defining reasonable proxies, Skalse et al. () introduce a notion of an “unhackable” proxy reward but find that a proxy is only unhackable if it is equivalent to the true reward up to scaling; this is clearly too restrictive of a definition.
The robust MDP literature has considered proxy rewards that differ from the true reward by at most some constant, e.g., $|R(s,a)-\tilde{R}(s,a)|\leq c$ for all $s,a$ (Derman et al., ). However, we find that in many realistic cases of reward hacking, there may be some states where a proxy differs from the true reward by an arbitrarily large amount, and so these definitions are generally not applicable. With regard to the second desideratum, defining a threshold of true reward, there is little existing literature.

## 4 Defining Reward Hacking

Despite the difficulties in formalizing reward hacking, we argue that both of our desiderata for a definition can be met by defining reward hacking with respect to a *base policy*. We show how this allows for both a precise definition of proxy rewards and a natural threshold for when reward hacking is occurring.

Characterizing proxy rewards.
To find a realistic definition of a good proxy reward, consider the process by which system designers create proxies. If they are hand-specified, we argue that designers usually imagine a distribution of behavior and design a reward that captures the objective under that distribution. For example, we study a traffic control simulator (Lopez et al., ; Vinitsky et al., ; Wu et al., ; Pan et al., ) where a small number of autonomous cars help to regulate traffic among a larger population of human drivers. In this case, a designer might choose the average speed of all vehicles as a proxy for improving traffic flow; this is a reasonable proxy because under the distribution of typical human driving behavior, higher speeds are associated with better traffic.

We can formalize this intuition by requiring that the proxy and true rewards be *correlated* over states and actions sampled from a *base policy*:

###### Definition 4.1 (Correlated proxy reward) .

An $r$-correlated proxy reward $\tilde{R}$ with respect to a base policy ${\pi_{\text{base}}}$ is one that has a correlation of $r>0$ with the true reward for state-action pairs sampled from the base policy:

$$ $\displaystyle\qquad\qquad\quad\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\left(\frac{\tilde{R}(s,a)-J({\pi_{\text{base}}},\tilde{R})}{\sigma_{\tilde{R}}}\right)\left(\frac{R(s,a)-J({\pi_{\text{base}}},R)}{\sigma_{R}}\right)\right]=r,$ $\displaystyle\textstyle\text{where}\quad\sigma_{\tilde{R}}=\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\big(\tilde{R}(s,a)-J({\pi_{\text{base}}},\tilde{R})\big)^{2}\right]\quad\text{and}\quad\sigma_{R}=\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\big(R(s,a)-J({\pi_{\text{base}}},R)\big)^{2}\right]$ $$

are the standard deviations of proxy and true rewards under the base policy.

This definition intuitively captures cases like the traffic environment, where we define the true reward function as the negative total commute time for all vehicles. Letting ${\pi_{\text{base}}}$ be an autonomous vehicle policy based on a common model of human driving behavior, we plot the distribution of true and proxy rewards in the top-left corner of Figure 2 and find that they are correlated: higher average speed tends to occur with lower commute times, and vice versa. This validates that average speed is a correlated proxy reward according to Definition 4.1.

Definition 4.1 also avoids too strongly constraining proxy rewards, unlike past formalisms. The proxy and true reward functions can diverge arbitrarily at some state-action pairs, as long as those state-action pairs have low or zero occupancy measure under the base policy.

Finally, Definition 4.1 captures *learned* proxy rewards in addition to hand-specified ones. Reward learning begins with collecting annotations—for example, preference comparisons (as in RLHF) or scalar feedback—for states and actions collected from a rollout policy—for example, the SFT model in RLHF. Then, a reward function is estimated via supervised learning over this dataset with an appropriate loss function. If the estimated reward function generalizes well in-distribution, then we should expect the true and learned rewards to be well correlated under states and actions sampled from the rollout policy, satisfying Definition 4.1.

Choosing a threshold for reward hacking.
As discussed in Section 3, the other difficulty in defining reward hacking is specifying a threshold of true reward below which performance is poor enough to be considered hacking. If we are already characterizing proxy rewards with respect to a base policy, then it makes sense to also use this base policy as a baseline for evaluating a policy that optimizes the proxy. If optimizing the proxy leads to worse true reward than the base policy achieves, then the system designer may as well simply use the base policy. On the other hand, if optimizing the proxy reward improves true reward over the base policy, then we can argue that optimizing the proxy succeeded.

###### Definition 4.2 (Reward hacking) .

Suppose a $\tilde{R}$ is an $r$-correlated proxy with respect to ${\pi_{\text{base}}}$ (Definition 4.1). Then we say *reward hacking* occurs when an optimal policy for $\tilde{R}$ has lower true reward than the base policy ${\pi_{\text{base}}}$:

$$ $J(\pi,R)<J({\pi_{\text{base}}},R)\qquad\text{for some}\qquad\pi\in\arg\max_{\pi}J(\pi,\tilde{R}).$ $$

For example, in the traffic control environment, optimizing for the proxy of average speed leads to the autonomous vehicles blocking a highway on-ramp. This decreases the speed of the cars on the ramp to zero, but allows cars on the highway to move at high speeds, overall increasing the proxy reward. However, the true reward becomes extremely low, since commute times for the cars on the on-ramp are aribitrarily long. This constitutes reward hacking under Definition 4.2 since the true reward of the proxy-optimized policy is lower than that achieved by typical human driving.

Verifying our definition experimentally.
To test whether Definition 4.2 accurately captures intuitive cases of reward hacking, we consider four realistic enviroments and an illustrative gridworld. We consider the aforementioned traffic simulator and two other environments originally studied by Pan et al. () as examples of reward hacking. SimGlucose (Man et al., ) is based on an FDA-approved simulator of Type 1 Diabetes patients in which a policy monitors glucose levels and administers insulin; the true reward captures patient health while the proxy reward prioritizes reducing the monetary cost of insulin. PandemicSimulator (Kompella et al., ) uses the SEIR model (Mwalili et al., ) to simulate the COVID-19 pandemic among a population; the policy controls the level of lockdown restrictions placed on the population by observing the results of testing. The proxy reward omits the political cost of certain decisions.

We also study RLHF, in which LLMs are optimized based on a reward function learned from human preferences. Following Gao et al. () and Coste et al. (), we use a large (more robust) reward model as the true reward function and a smaller (less robust) one as the proxy. Finally, as an interpretable and illustrative example, we include the tomato-watering gridworld from Leike et al. (). In this environment, a robot that waters tomatoes receives true reward depending on how many plants are watered; however, the proxy additionally rewards the robot for standing in a sprinkler where it appears all the tomatoes are watered.
Besides the gridworld, all of our environments reflect complex, realistic tasks with large or infinite state state spaces.
To test our definition of reward hacking, for each of the five environments we construct a natural base policy. See Appendix C for more details about the environments and base policies.

The top row of Figure 2 shows the distribution of true and proxy reward values for state-action pairs sampled from these base policies in each environment. We find that in all cases, the proxy rewards correlate with the true reward, satisfying our definition of a correlated proxy. The middle row of Figure 2 shows the result of optimizing for the proxies: for each environment, we plot the distribution of true and proxy reward values for state-action pairs reached by a policy optimized on the proxy reward via RL. In all cases, we see that the true reward drops significantly compared to the base policy. Thus, these intuitive cases of reward hacking are captured by Definition 4.2.

Figure: Figure 3: Unlike RLHF, which regularizes *action distribution* (AD) divergence to prevent reward hacking, our results suggest regularizing using *occupancy measure* (OM) divergence. These plots of the glucose monitoring environment show the typical ADs and OMs of two policies. $\pi$ is close to ${\pi_{\text{base}}}$ in AD; it is more likely to give slightly less insulin because the proxy reward penalizes monetary cost. However, this leads to a vastly different OM, with typical glucose levels for $\pi$ far outside the healthy range (dotted lines). Thus, regularizing ADs to be close to ${\pi_{\text{base}}}$ is not enough to prevent reward hacking; instead, divergence between the OMs better captures the reward hacking behavior.

###### Definition 4.1 (Correlated proxy reward) .

###### Definition 4.2 (Reward hacking) .

## 5 Mitigating Reward Hacking with Occupancy Measure Regularization

We now discuss how our new definition of reward hacking can lead to better methods for preventing it. Ideally, we would like to be able to optimize a proxy reward function and have it translate into an improvement in true reward over the base policy. To understand how this might be possible, consider again the traffic environment. The reward hacking policy exhibits behavior which is very unlikely under the base policy: human drivers hardly ever stop on on-ramps and refuse to move. That is, optimizing the proxy reward leads to *out-of-distribution states and actions* where the correlation that made the proxy good in the first place breaks down. Visually, this can be seen in the second row of Figure 2: the reward hacking policy usually finds state-action pairs with high proxy reward and low true reward that were very unlikely to be reached by the base policy.

Thus, to prevent reward hacking, one solution could be to optimize the proxy reward while avoiding states that are unlikely under the base policy. The following theorem formalizes this idea.

###### Theorem 5.1 .

Suppose that $\tilde{R}$ is an $r$-correlated-proxy for the true reward function $R$, and let $\sigma_{\tilde{R}}$ and $\sigma_{R}$ be defined as in Definition 4.1.
Then for any policy $\pi$ such that $\mu_{\pi}\ll\mu_{{\pi_{\text{base}}}}$ (i.e., $\mu_{{\pi_{\text{base}}}}(s,a)=0\Rightarrow\mu_{\pi}(s,a)=0$), we have

$$ $\frac{J(\pi,R)-J({\pi_{\text{base}}},R)}{\sigma_{R}}\geq\frac{1}{r}\left(\frac{J(\pi,\tilde{R})-J({\pi_{\text{base}}},\tilde{R})}{\sigma_{\tilde{R}}}-\sqrt{(1-r^{2})\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}\right),$ (4) $$

where $\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)=\mathbb{E}_{\mu_{\pi}}\left[\frac{\mu_{\pi}(s,a)}{\mu_{\pi_{\text{base}}}(s,a)}-1\right]$ is the $\chi^{2}$ divergence between $\mu_{\pi}$ and $\mu_{\pi_{\text{base}}}$.

See Appendix A for the proof.
Equation (4) gives a lower bound on how much the policy $\pi$ improves over the base policy ${\pi_{\text{base}}}$ on the true reward, normalized by the standard deviation of the true reward under the base policy. This is exactly we would like to maximize, but we can’t optimize the unknown true reward directly, so it makes sense to instead optimize the right hand side of (4). The RHS consists of two terms. The first is the improvement of the policy $\pi$ over the base policy under the proxy reward, again normalized; optimizing this term alone often leads to reward hacking. However, the second term penalizes the divergence of $\pi$’s occupancy measure from that of ${\pi_{\text{base}}}$. By incentivizing $\pi$ to achieve high proxy reward but also stay close to the base policy, $\pi$ can achieve high true reward.

Scaling the RHS of (4) and removing terms that are constant in $\pi$ suggests using the following regularized policy optimization objective to avoid reward hacking:

$$ $\text{maximize}\quad J(\pi,\tilde{R})-\lambda\sqrt{\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}\qquad\text{where}\qquad\lambda=\sigma_{\tilde{R}}\sqrt{1-r^{2}}.$ (5) $$

The amount of regularization needed to improve on the base policy depends on the strength of correlation $r$. The higher the correlation, the lower the regularization strength $\sqrt{1-r^{2}}$. Given the prefactor of $\frac{1}{r}$ on the RHS of (4), it may appear that a lower correlation leads to a larger gain in true reward. However, Lemma A.1 shows that in fact, the lower bound decreases as a function of $r$.

Comparison to KL regularization in RLHF.
Regularization to a base policy is already widely used to prevent reward hacking in RLHF (Stiennon et al., ; Bai et al., ). Specifically, a KL penalty is applied between the distributions of responses generated by the optimized policy and the SFT policy. In our setting, we can write this as

$$ $\textstyle\text{maximize}\quad J(\pi,\tilde{R})-\lambda(1-\gamma)\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}D_{\text{KL}}\big(\pi(\cdot\mid s_{t})\;\|\;\mu_{\pi_{\text{base}}}(\cdot\mid s_{t})\big)\right].$ (6) $$

That is, in RLHF, the expected KL divergence between the action distributions of $\pi$ and ${\pi_{\text{base}}}$ is penalized. Action distribution divergence is easy to calculate and optimize, and training LLMs with (6) seems to work well in practice. However, unlike our objective in (5), (6) lacks theoretical guarantees, and it is unclear if it works in other environments.

Our regularized objective in (5) differs in two main ways from the KL regularization used in RLHF. First, our results suggest optimizing the *occupancy measure* (OM) divergence between policies, whereas RLHF uses the *action distribution* (AD) divergence. Second, Theorem 5.1 applies to $\chi^{2}$ divergence, while RLHF uses KL divergence. In the remainder of this section, we explore intuitively why OM divergence is preferable to AD divergence, and why $\chi^{2}$ divergence is superior to KL divergence. Then, in Section 6, we empirically explore applying different types of regularization to prevent reward hacking in five environments.

Figure: Figure 4: Our theory also suggests preventing reward hacking by regularizing with *$\chi^{2}$ divergence* instead of *KL divergence* like prior work. This plot illustrates how $\chi^{2}$ regularization is more effective at preventing reward hacking in RLHF. Both divergences can be written as the expectation of a function $g\big(\log\big(\mu_{\pi}(s,a)/\mu_{{\pi_{\text{base}}}}(s,a)\big)\big)$ which penalizes state-action pairs more the further the log-ratio is from zero. The $g(\cdot)$ associated with KL divergence only increases slowly for large log-ratios, so policies trained with KL divergence may produce nonsensical text. In contrast, the $g(\cdot)$ for $\chi^{2}$ divergence increases exponentially, better constraining the LLM to produce text similar to the SFT policy.

Occupancy measure vs. action distribution regularization.
While AD regularization works well for RLHF, this may be because RLHF is essentially a contextual bandit problem, meaning that $\gamma=0$; in this case, OM and AD divergence are equivalent (see Appendix A.3). However, in other cases, AD divergence may not suffice to prevent reward hacking behavior. This is because, in longer-horizon environments, a small change in action distribution at a single state can lead to a much higher probability of reaching undesirable states.
Figure 3 shows an example of this in the glucose monitoring environment: policies that are close in action distribution regularization produce vastly different patient glucose levels. Occupancy measure regularization avoids this issue by directly preventing the distribution of glucose levels from differing too much from the base policy.

$\chi^{2}$ vs. KL divergence.
Compared to KL divergence, $\chi^{2}$ divergence may be superior for preventing reward hacking because it more strongly penalizes out-of-distribution state-action pairs.
To illustrate this, we can write both divergences as expectations over functions of the log-ratio of the occupancy measures, which we denote $d(s,a)$:

$$ $\textstyle D_{\text{KL}}(\mu_{\pi}\;\|\;\mu_{{\pi_{\text{base}}}})=\mathbb{E}_{\mu_{\pi}}\left[d(s,a)+e^{-d(s,a)}\right]\qquad\chi^{2}(\mu_{\pi}\;\|\;\mu_{{\pi_{\text{base}}}})=\mathbb{E}_{\mu_{\pi}}\left[e^{d(s,a)}+e^{-d(s,a)}\right]$ $$

$$ $\text{where}\qquad d(s,a)=\log\big(\mu_{\pi}(s,a)/\mu_{{\pi_{\text{base}}}}(s,a)\big).$ (7) $$

As $d(s,a)$ increases, the optimized policy is visiting state-action pairs that are less likely under the base policy. However, KL divergence only penalizes increases in $d(s,a)$ linearly, while $\chi^{2}$ penalizes them exponentially, resulting in stronger regularization even with a lower coefficient. Figure 4 plots the functions in (7) and shows how in practice $\chi^{2}$ divergence better prevents reward hacking in RLHF.

**Table 1: We compare using various types of regularization to prevent reward hacking in the five environments form Figure 2. The median true reward and standard deviation across 5 random seeds is shown for the best regularization coefficient for each type of regularization. The bottom rows show results for the baselines: the base policy ${\pi_{\text{base}}}$, a policy trained on the proxy reward without regularization (exhibiting reward hacking), and a policy trained on the true reward function (impossible in practice, but included as an upper bound on performance). We find that occupancy measure regularization consistently improves on action distribution regularization, and that $\chi^{2}$ divergence is often superior to KL divergence.**
|  | Environment |  |  |  |  |
| --- | --- | --- | --- | --- | --- |
|  | Traffic | Pandemic | Glucose | RLHF | AI safety |
| Method | control | mitigation | monitoring |  | gridworld |
|  | $(\times 10^{3})$ |  | $(\times 10^{3})$ |  |  |
| Action dist. $\chi^{2}$ | -1.29 $\pm$ 0.10 | -12.29 $\pm$ 0.05 | -74.8 $\pm$ 11.8 | 16.94 $\pm$ 0.07 | 6.24 $\pm$ 0.09 |
| State occupancy $\chi^{2}$ | -2.18 $\pm$ 0.38 | -10.68 $\pm$ 0.15 | -54.7 $\pm$ 1.0 | — | 9.07 $\pm$ 0.06 |
| State-action occupancy $\chi^{2}$ | -1.15 $\pm$ 0.05 | -11.17 $\pm$ 0.17 | -47.6 $\pm$ 0.6 | — | 9.17 $\pm$ 0.11 |
| Action dist. KL | -1.33 $\pm$ 0.05 | -12.20 $\pm$ 0.06 | -73.4 $\pm$ 8.3 | 16.81 $\pm$ 0.27 | 6.33 $\pm$ 0.11 |
| State occupancy KL | -1.34 $\pm$ 22.6 | -10.24 $\pm$ 0.54 | -58.4 $\pm$ 3.4 | — | 7.07 $\pm$ 0.11 |
| State-action occupancy KL | -1.25 $\pm$ 0.06 | -11.73 $\pm$ 0.19 | -48.9 $\pm$ 0.5 | — | 6.86 $\pm$ 0.17 |
| ${\pi_{\text{base}}}$ | -2.28 $\pm$ 0.00 | -12.26 $\pm$ 0.00 | -72.6 $\pm$ 0.0 | 16.37 $\pm$ 0.00 | 5.86 $\pm$ 0.00 |
| No regularization | -57.38 $\pm$ 3.53 | -29.57 $\pm$ 6.86 | -599.0 $\pm$ 1.6 | 9.16 $\pm$ 0.80 | 2.35 $\pm$ 0.14 |
| Training with true reward | -0.93 $\pm$ 0.11 | -2.65 $\pm$ 0.83 | -43.4 $\pm$ 0.8 | — | 8.54 $\pm$ 0.12 |

###### Theorem 5.1 .

## 6 Experiments

We now show that our theoretical results—which suggest $\chi^{2}$ occupancy measure regularization can prevent reward hacking—lead to empirical success in realistic environments.

Practical occupancy measure regularization.
Occupancy measure regularization is more difficult to implement in practice compared to action distribution regularization. While AD regularization can be added as a loss term to deep RL algorithms like proximal policy optimization (PPO), OM divergences cannot be calculated in closed form. Instead, we follow several previous works (e.g., Ho & Ermon ; Kang et al. ) and use a discriminator network to approximate OM divergences. Specificaly, in Appendix B, we show the objective in (5) can be optimized via policy gradient with an adjusted reward function that depends on a discriminator $\hat{d}_{\phi}$:

$$ $\displaystyle\textstyle R^{\prime}(s,a)$ $\displaystyle\textstyle=\tilde{R}(s,a)-\frac{\lambda}{\sqrt{\widehat{\chi^{2}}}}e^{\hat{d}_{\phi}(s,a)}\qquad\text{where}\qquad\widehat{\chi^{2}}=\mathbb{E}_{\mu_{\pi}}\Big[\,e^{\hat{d}_{\phi}(s_{t},a_{t})}-1\,\Big]$ $\displaystyle\text{and}\qquad\phi$ $\displaystyle=\arg\min_{\phi}\;\mathbb{E}_{\mu_{\pi}}\Big[\,\log(1+e^{-\hat{d}_{\phi}(s,a)})\,\Big]+\,\mathbb{E}_{\mu_{\pi_{\text{base}}}}\Big[\,\log(1+e^{\hat{d}_{\phi}(s,a)})\,\Big].$ (8) $$

That is, optimizing a discriminator network $\hat{d}_{\phi}$ to minimize the given loss can be used to estimate and optimize $\chi^{2}$ OM divergence.
We alternately train $\hat{d}_{\phi}$ via gradient descent on (8) and the policy $\pi$ via PPO based on the adjusted reward. We call this algorithm Occupancy-Regularized Policy Optimization (ORPO).
See Appendix B for a full derivation of the approximations used in ORPO and Algorithm 1 for a formal description. We use a similar strategy for regularizing based on OM KL divergence; see the appendix for details.

Figure: Figure 5: The true reward achieved by policies regularized with varying amounts of action distribution or occupancy measure regularization using $\chi^{2}$ and KL divergence. The x-axis is the regularization coefficient $\lambda$ normalized by the standard deviation of proxy rewards under the base policy. Dots indicate the median reward and the shaded area is the range over random seeds. For RLHF, AD and OM regularization are equivalent, which is why OM regularization results are not shown for that column.

Experimental setup.
In each of the five environments shown in Figure 2, we train policies with four types of regularization towards the base policy: AD KL, AD $\chi^{2}$, OM KL, and OM $\chi^{2}$. In Appendix A.4, we show that Theorem 5.1 also holds for state-only occupancy measures if the environment’s reward function does not depend on the action; thus, we experiment with regularizing based on both state-action and state-only OM divergence. For each environment and type of regularization, we test a number of regularization coefficients $\lambda$. Theorem 5.1 suggests setting $\lambda=\sigma_{\tilde{R}}\sqrt{1-r^{2}}$ for an $r$-correlated proxy, so for $\chi^{2}$ regularization we test a range of values $\lambda=c\,\sigma_{\tilde{R}}$ from $c=1$ to $10^{-2}$ ($10^{-4}$ for RLHF). Since it is less clear theoretically how to set the coefficient $\lambda$ for KL regularization, we experiment with a wider range of values.

We train each combination of {$\chi^{2}$ divergence, KL divergence} $\times$ {AD, state OM, state-action OM} $\times$ $\{\lambda_{1},\lambda_{2},\dots\}$ with five random seeds and measure the resulting policies’ expected returns under the true reward. As a baseline, we train a policy without any regularization, which leads to reward hacking in all environments. We also train a policy directly on the true reward function as an upper limit for performance. In RLHF, we only consider AD regularization since it is equivalent to OM regularization for LLM chatbots (see Appendix A.3). We do not train a policy on the true reward for RLHF as we found it could be hacked with enough optimization pressure. See Appendix D for all hyperparameter and experiment details.

Results.
The results of our experiments are shown in Table 1 and Figure 5. Table 1 shows the median true reward with the best coefficient for each type of regularization. We find that OM regularization consistently outperforms AD regularization across the four non-RLHF environments. In two environments (glucose monitoring and pandemic mitigation), AD regularization fails to improve on the base policy’s true reward at all. Furthermore, $\chi^{2}$ regularization tends to perform similarly to or better than KL regularization across all environments. In RLHF in particular, $\chi^{2}$ regularization leads to a larger improvement over the base policy compared to the industry-standard KL penalty, and is more stable across seeds as well.

However, comparing the types of regularization based only on the best coefficient is unrealistic, since in practice designers must choose a coefficient $\lambda$ without access to the true reward. In Figure 5, we show the true reward achieved when training with each type of regularization across a range of $\lambda$ values. In addition to performing best with an optimal coefficient, $\chi^{2}$ and OM regularization seem to also perform well over a larger range of coefficients compared to AD regularization.

In Appendix E, we present the full results of our experiments as well as ablations of ORPO.

## 7 Conclusion

We have introduced a new definition for reward hacking based on correlation between a proxy reward function and the unknown true reward that breaks down when optimizing the proxy. Furthermore, we leveraged this definition to show theoretically and empirically that $\chi^{2}$ occupancy measure regularization can effectively prevent reward hacking. Our results have implications for settings, like RLHF, where RL is used to optimize complex, hard-to-specify objectives. We suggest that the heuristic KL penalty used currently should be replaced by a more principled form of regularization. While OM and AD regularization are equivalent for today’s formulation of RLHF as a contextual bandit, they will no longer remain so as LLM-based agents are optimized over multi-turn conversations or with tool use (Wang et al., ; Abdulhai et al., ; Shani et al., ). Thus, our results provide a principled path to continuing to ensure the safety of increasingly powerful AI systems.

#### Acknowledgments

We would like to thank Sam Toyer, Micah Carroll, Katie Kang, Dibya Ghosh, and Cam Allen for helpful discussions and feedback on drafts.

This work was supported by the Semiconductor Research Corporation (SRC)’s Center for the Co-Design of Cognitive Systems (COCOSYS) as well as the NSF’s Human Centered Computing (HCC) program. Cassidy Laidlaw is supported by an Open Philanthropy AI Fellowship and a National Defense Science and Engineering Graduate (NDSEG) Fellowship.

## Appendix A Proofs

### A.1 Proof of Theorem 5.1

See 5.1

###### Proof.

For simplicity of exposition, define

$$ $\displaystyle Z(s,a)$ $\displaystyle=\frac{R(s,a)-J({\pi_{\text{base}}},R)}{\sigma_{R}}\qquad\text{and}\qquad\tilde{Z}(s,a)=\frac{\tilde{R}(s,a)-J({\pi_{\text{base}}},\tilde{R})}{\sigma_{\tilde{R}}}.$ $$

Using (3), we can rewrite (4) as

$$ $\displaystyle\mathbb{E}_{\mu_{\pi}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]$ $\displaystyle\leq\sqrt{(1-r^{2})\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}.$ $$

Then, the left hand side can be rewritten as

$$ $\displaystyle\mathbb{E}_{\mu_{\pi}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]$ $\displaystyle\quad=\mathbb{E}_{\mu_{\pi}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]-\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]+\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]$ $\displaystyle\quad=\mathbb{E}_{\mu_{\pi}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]-\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right],$ $$

since by definition $\mathbb{E}_{\mu_{\pi_{\text{base}}}}[\tilde{Z}(s,a)]=\mathbb{E}_{\mu_{\pi_{\text{base}}}}[Z(s,a)]=0$. Applying the Cauchy-Schwartz inequality to this difference gives

$$ $\displaystyle\mathbb{E}_{\mu_{\pi}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]-\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]$ $\displaystyle\quad=\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]\Big[\mu_{\pi}(s,a)-\mu_{\pi_{\text{base}}}(s,a)\Big]$ $\displaystyle\quad=\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\left[\sqrt{\mu_{\pi_{\text{base}}}(s,a)}\Big(\tilde{Z}(s,a)-r\,Z(s,a)\Big)\right]\left[\frac{\mu_{\pi}(s,a)-\mu_{\pi_{\text{base}}}(s,a)}{\sqrt{\mu_{\pi_{\text{base}}}(s,a)}}\right]$ $\displaystyle\quad\leq\sqrt{\left(\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\mu_{\pi_{\text{base}}}(s,a)\left[\tilde{Z}(s,a)-r\,Z(s,a)\right]^{2}\right)\left(\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\frac{(\mu_{\pi}(s,a)-\mu_{\pi_{\text{base}}}(s,a))^{2}}{\mu_{\pi_{\text{base}}}(s,a)}\right)}$ $\displaystyle\quad=\sqrt{\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\Big(\tilde{Z}(s,a)-r\,Z(s,a)\Big)^{2}\right]\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}.$ (9) $$

The expectation can be calculated as

$$ $\displaystyle\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\Big(\tilde{Z}(s,a)-r\,Z(s,a)\Big)^{2}\right]$ $\displaystyle\quad=\text{Var}_{\mu_{\pi_{\text{base}}}}\left(\tilde{Z}(s,a)-r\,Z(s,a)\right)$ $\displaystyle\quad=\text{Var}_{\mu_{\pi_{\text{base}}}}\left(\tilde{Z}(s,a)\right)+r^{2}\,\text{Var}_{\mu_{\pi_{\text{base}}}}\left(Z(s,a)\right)-2r\,\text{Cov}_{\mu_{\pi_{\text{base}}}}\left(\tilde{Z}(s,a),Z(s,a)\right)$ $\displaystyle\quad=1-r^{2},$ $$

using the fact that both reward functions have unit variance under the base policy and that their correlation is $r$. Plugging this into (9) gives the desired result.
∎

###### Proof.

### A.2 Understanding the Lower Bound in Theorem 5.1

Denote by

$$ $L(\pi)=\frac{1}{r}\left(\frac{J(\pi,\tilde{R})-J({\pi_{\text{base}}},\tilde{R})}{\sigma_{\tilde{R}}}-\sqrt{(1-r^{2})\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}\right)$ $$

the lower bound on increase in the true reward which is the RHS of (4). One surprising observation is that this lower bound seems to be *increasing* as the proxy becomes *less* correlated with the true reward. This would suggest that a less correlated proxy leads to better optimization of the true reward. However, as the following lemma shows, $L(\pi)$ is actually decreasing in $r$.

###### Lemma A.1 .

Under the same conditions as Theorem 5.1, the lower bound $L(\pi)$ satisfies

$$ $L(\pi)\leq\frac{1-\sqrt{1-r^{2}}}{r}\sqrt{\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}.$ (10) $$

This shows that the lower bound can be at most a factor of $\frac{1-\sqrt{1-r^{2}}}{r}$ times the divergence between the learned and base policies’ occupancy measures. This factor is increasing in $r$ and asymptotes to $r/2$ as $r\to 0$.

###### Proof.

Using the same notation as the proof of Theorem 5.1, we can rewrite the lower bound as

$$ $\displaystyle L(\pi)=\frac{1}{r}\left(\mathbb{E}_{\pi}\left[\tilde{Z}(s,a)\right]-\sqrt{(1-r^{2})\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}\right).$ $$

Following a similar argument to that in the proof of Theorem 5.1, we can write

$$ $\displaystyle\mathbb{E}_{\pi}\left[\tilde{Z}(s,a)\right]$ $\displaystyle=\mathbb{E}_{\pi}\left[\tilde{Z}(s,a)\right]-\mathbb{E}_{\pi_{\text{base}}}\left[\tilde{Z}(s,a)\right]$ $\displaystyle=\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\left[\sqrt{\mu_{\pi_{\text{base}}}(s,a)}\tilde{Z}(s,a)\right]\left[\frac{\mu_{\pi}(s,a)-\mu_{\pi_{\text{base}}}(s,a)}{\sqrt{\mu_{\pi_{\text{base}}}(s,a)}}\right]$ $\displaystyle\leq\sqrt{\left(\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\mu_{\pi_{\text{base}}}(s,a)\tilde{Z}(s,a)^{2}\right)\left(\sum_{(s,a)\in\mathcal{S}\times\mathcal{A}}\frac{(\mu_{\pi}(s,a)-\mu_{\pi_{\text{base}}}(s,a))^{2}}{\mu_{\pi_{\text{base}}}(s,a)}\right)}$ $\displaystyle=\sqrt{\mathbb{E}_{\mu_{\pi_{\text{base}}}}\left[\tilde{Z}(s,a)^{2}\right]\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}$ $\displaystyle=\sqrt{\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}.$ $$

Combining this with the definition of $L(\pi)$ gives

$$ $\displaystyle L(\pi)$ $\displaystyle\leq\frac{1}{r}\left(\sqrt{\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}-\sqrt{(1-r^{2})\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)}\right)$ $\displaystyle=\frac{1-\sqrt{1-r^{2}}}{r}\sqrt{\chi^{2}\left(\mu_{\pi}\|\mu_{\pi_{\text{base}}}\right)},$ $$

which completes the proof.
∎

###### Lemma A.1 .

###### Proof.

### A.3 Action Distribution and Occupancy Measure Divergences in Contextual Bandits

In our setting, a *contextual bandit* can be defined as an MDP with $\gamma=0$; then, the return of the policy $\pi$ under a reward function $R$ is given by

$$ $J(\pi,R)=\mathbb{E}_{s\sim\mu_{0}(\cdot),a\sim\pi(\cdot\mid s)}[R(s,a)].$ $$

That is, a single state is sampled from the initial state distribution $\mu_{0}$, and then a single action is sampled from the policy $\pi$ conditioned on that state. For example, RLHF for LLMs generally follows this setting: a prompt is sampled from a dataset, the LLM generates a response, and the reward is calculated based on the prompt and response.

In this setting, it is simple to show that the action distribution and occupancy measure divergences are equivalent.

###### Lemma A.2 .

Let $D_{f}(P\;\|\;Q)=\mathbb{E}_{x\sim Q}\left[f\big(P(x)/Q(x)\big)\right]$ be the $f$-divergence between two distributions $P$ and $Q$. Then, for any two policies $\pi,\pi^{\prime}$ in a contextual bandit, we have

$$ $D_{f}(\mu_{\pi}\;\|\;\mu_{\pi^{\prime}})=\mathbb{E}_{s\sim\mu_{0}(\cdot)}\left[D_{f}(\pi\;\|\;\pi^{\prime})\right].$ $$

Lemma A.2 applies to any $f$-divergence, including the KL and $\chi^{2}$ divergences we study in this paper.

###### Proof.

We have

$$ $\displaystyle D_{f}(\mu_{\pi}\;\|\;\mu_{\pi^{\prime}})$ $\displaystyle=\mathbb{E}_{s\sim\mu_{0}(\cdot),a\sim\pi^{\prime}(\cdot\mid s)}\left[f\left(\frac{\mu_{\pi}(s,a)}{\mu_{\pi^{\prime}}(a\mid s)}\right)\right]$ $\displaystyle=\mathbb{E}_{s\sim\mu_{0}(\cdot)}\left[\mathbb{E}_{a\sim\pi^{\prime}(\cdot\mid s)}\left[f\left(\frac{\mu_{0}(s)\pi(a\mid s)}{\mu_{0}(s)\pi^{\prime}(a\mid s)}\right)\right]\right]$ $\displaystyle=\mathbb{E}_{s\sim\mu_{0}(\cdot)}\left[D_{f}(\pi\;\|\;\pi^{\prime})\right],$ $$

which completes the proof.
∎

###### Lemma A.2 .

###### Proof.

### A.4 State-only occupancy measures

We can also consider state-only occupancy measures, defined as

$$ $\mu_{\pi}(s)=(1-\gamma)\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}\mathds{1}\{s_{t}=s\}\right].$ $$

In many environments, the reward functions only depend on the state, i.e., $R(s,a)=R(s)$ and $\tilde{R}(s,a)=\tilde{R}(s)$. In this case, Theorem 5.1 holds for state-only occupancy measures as well. The proof is identical to the proof of Theorem 5.1, but replacing expectations and sums over state-action pairs with expectations over states.

## Appendix B Derivation of Occupancy-Regularized Policy Optimization

In this appendix section, we show how to derive the approximations used for Occupancy-Regularized Policy Optimization (ORPO). As a reminder, we would like to optimize

$$ $J(\pi_{\theta},\tilde{R})-\lambda\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}.$ $$

We can rewrite its gradient as

$$ $\displaystyle\nabla_{\theta}\Big(J(\pi_{\theta},\tilde{R})-\lambda\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}\Big)$ $\displaystyle\quad=\nabla_{\theta}J(\pi_{\theta},\tilde{R})-\frac{\lambda\nabla_{\theta}\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}{2\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}}$ $\displaystyle\quad=\nabla_{\theta}\left(\sum_{s,a}\mu_{\pi_{\theta}}(s,a)\tilde{R}(s,a)\right)-\frac{\lambda}{2\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}}\nabla_{\theta}\left(\sum_{s,a}\frac{\mu_{\pi_{\theta}}(s,a)^{2}}{\mu_{{\pi_{\text{base}}}}(s,a)}-1\right)$ $\displaystyle\quad=\sum_{s,a}\left[\nabla_{\theta}\mu_{\pi_{\theta}}(s,a)\tilde{R}(s,a)-\frac{\lambda}{2\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}}\nabla_{\theta}\left(\frac{\mu_{\pi_{\theta}}(s,a)^{2}}{\mu_{{\pi_{\text{base}}}}(s,a)}-1\right)\right]$ $\displaystyle\quad=\sum_{s,a}\left[\nabla_{\theta}\mu_{\pi_{\theta}}(s,a)\tilde{R}(s,a)-\frac{\lambda}{2\sqrt{\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)}}\frac{2\mu_{\pi_{\theta}}(s,a)}{\mu_{{\pi_{\text{base}}}}(s,a)}\nabla_{\theta}\mu_{\pi_{\theta}}(s,a)\right]$ $\displaystyle\quad=\sum_{s,a}\Big(\nabla_{\theta}\mu_{\pi_{\theta}}(s,a)\Big)\left(\tilde{R}(s,a)-\frac{\lambda}{\sqrt{\chi^{2}(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}})}}\frac{\mu_{\pi_{\theta}}(s,a)}{\mu_{\pi_{\text{base}}}(s,a)}\right).$ $$

As described in the main text, policy gradient algorithms can approximate this type of gradient by using an augmented reward function

$$ $R^{\prime}(s,a)=\tilde{R}(s,a)-\frac{\lambda}{\sqrt{\chi^{2}(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}})}}\frac{\mu_{\pi_{\theta}}(s,a)}{\mu_{\pi_{\text{base}}}(s,a)}.$ (11) $$

However, two terms in (11) cannot be computed directly: the current $\chi^{2}$ divergence between occupancy measures, and the ratio of the occupancy measures $\mu_{\pi_{\theta}}(s,a)/\mu_{\pi_{\text{base}}}(s,a)$. We first show how to approximate the latter using a discriminator network $\hat{d}^{\phi}(s,a)$, trained to optimize

$$ $\phi=\arg\min_{\phi}\;\mathbb{E}_{\mu_{\pi_{\theta}}}\Big[\,\log(1+e^{-\hat{d}_{\phi}(s,a)})\,\Big]\\ \textstyle+\,\mathbb{E}_{\mu_{\pi_{\text{base}}}}\Big[\,\log(1+e^{\hat{d}_{\phi}(s,a)})\,\Big].$ (12) $$

It is well-known that exactly optimizing the loss function in (12) gives

$$ $\hat{d}^{\phi}(s,a)=\log\frac{\mu_{\pi_{\theta}}(s,a)}{\mu_{\pi_{\text{base}}}(s,a)}.$ (13) $$

Furthermore, the OM $\chi^{2}$ divergence is given by

$$ $\chi^{2}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right)=\mathbb{E}_{\mu_{\pi_{\theta}}}\left[\frac{\mu_{\pi_{\theta}}(s,a)}{\mu_{\pi_{\text{base}}}(s,a)}-1\right]=\mathbb{E}_{\mu_{\pi_{\theta}}}\left[e^{\hat{d}^{\phi}(s,a)}-1\right]\eqqcolon\widehat{\chi^{2}}.$ (14) $$

Combining (13) and (14) shows that the augmented reward in (11) can be rewritten as

$$ $\textstyle R^{\prime}(s,a)=\tilde{R}(s,a)-\frac{\lambda}{\sqrt{\widehat{\chi^{2}}}}e^{\hat{d}_{\phi}(s,a)}.$ $$

Putting all the steps together, the following algorithm formalizes ORPO:

Figure: Algorithm 1 Occupancy-Regularized Policy Optimization (ORPO).

A similar approach can also be used to optimize the proxy reward regularized by KL divergence

$$ $J(\pi_{\theta},\tilde{R})-\lambda\,D_{\text{KL}}\left(\mu_{\pi_{\theta}}\|\mu_{\pi_{\text{base}}}\right),$ $$

by changing the augmented reward in Line 6 of Algorithm 1 to

$$ $R^{\prime}(s,a)=\tilde{R}(s,a)-\lambda\,\hat{d}_{\phi}(s,a).$ $$

## Appendix C Environment Details

Here, we discuss the details of the five reward-hacking environments we study.

### C.1 Traffic control

The traffic control environment, based on the Flow simulator (Wu et al., ), simulates a group of human-controlled and RL-controlled vehicles on an on-ramp attempting to merge into traffic on a highway. The true reward prioritizes a small mean commute time, while the proxy reward is the average velocity of all cars. When reward hacking, the RL controlled vehicle on the on-ramp stops indefinitely and lets cars continue forward at high speeds on the highway, which maximizes the proxy reward but increases the commute times of cars on the on-ramp infinitely.
As the base policy for the traffic environment we used the Intelligent Driver Model (IDM), a standard approximation of human driving behavior (Treiber et al., ). In practice, base policies are often learned via imitation learning, so to simulate this we generate data from the IDM controller and train a behavioral cloning (BC) policy using the generated data.

Here, the green cars are controlled by the human driver model IDM controller, and the blue cars are controlled by RL:

This particular frame showcases reward hacking behavior. The blue RL-controlled vehicle has stopped completely on the on-ramp, blocking cars behind it. This increases the average velocity of all vehicles in the simulation, as the cars on the straightway are able to continue speeding along the road without having to wait for merging cars. However, the true reward (negative average commute time) decreases endlessly as the cars on the on-ramp wait.

### C.2 Glucose monitoring

The SimGlucose blood glucose monitoring environment is an extension of the FDA-approved glucose monitoring simulator proposed by Man et al. () for Type 1 Diabetes patients (Fox et al., ):

The RL agent (bottom) controls the insulin administered to a simulated patient in order to maintain healthy glucose levels. The true reward is a standard measure of health risk for the patient, but the proxy reward is misaligned and prioritizes the monetary cost of insulin. As the safe baseline policy, we train a BC policy based on data generated by a PID controller with parameters tuned by the original designers of the simulator (Steil, ).

### C.3 Pandemic mitigation

PandemicSimulator simulates a population’s infection dynamics using the SEIR model (Mwalili et al., ; Kompella et al., ):

The RL agent chooses the level of lockdown restrictions placed on the population by observing the results of testing.
The proxy reward function omits the political cost associated with certain decisions.
Our base policy is trained via BC on a combination of hand-specified and real-world strategies employed by governments during the pandemic, which were also used by Kompella et al. () as baselines.

### C.4 RLHF

We base our RLHF environment on the work of Coste et al. (), who study overoptimization of LLM-based reward models. The proxy reward model we use is fine-tuned from Pythia-70M (Biderman et al., ) on the AlpacaFarm (Dubois et al., ) preference dataset. As a true reward model, we use the Llama 3 Tulu V2 8B RM from AI2 (Ivison et al., ). As a base policy, we use Coste et al.’s SFT policy, which was fine-tuned from Python-1.4B on the AlpacaFarm SFT data. We evaluate policies’ true and proxy rewards on responses sampled for a held out set of prompts.

### C.5 Tomato-watering gridworld

The tomato environment contains a sprinkler state where the agent perceives all tomatoes as being watered and thus receives high proxy reward but no true reward. We train a base policy using the true reward function, and then add a 10% chance of taking a random action to ensure there is room to improve upon it.

The gray squares in the environment represent walls, and the white squares represent open spaces where the agent can travel:

The sprinkler state is down a narrow hallway, and on the other end a tomato is down another narrow hallway.

## Appendix D Experiment Details

### D.1 Non-LLM experiments

We implement ORPO using RLLib (Liang et al., ) and PyTorch (Paszke et al., ).

##### Network architectures

For the pandemic, traffic, and tomato-watering environments, we use policy networks based on fully-connected networks with 2 layers of 128 units, 4 layers of 512 units, and 4 units of 512 units, respectively. The policy model for the glucose environment is a basic LSTM network with 3 layers of 64 units each. We made this choice since the observation of the environment contains continuous historical information about the patient’s blood glucose levels and previously administered insulin.

The discriminator model for all four non-LLM environments is a fully connected network with 2 layers of 256 units each. We found that the discriminator architecture did not need to be tuned to each environment.

##### Policy initialization

Initializing using an imitation learning policy has been shown to effectively speed up the learning process (Laidlaw et al., ; Uchendu et al., ) and is used in practice for RLHF (Stiennon et al., ), so we initialize our policies using the specified ${\pi_{\text{base}}}$ for the more realistic traffic, glucose, and pandemic environments.

##### Hyperparameters

Some hyperparameters for the traffic environment were tuned by Pan et al. (). We primarily tuned the hyperparameters listed below in order to ensure that the proxy reward would be properly optimized and reward hacking would occur without regularization. This enables to see if the various regularization methods actually succeed at preventing reward hacking.

**Table 2: PPO/ORPO hyperparameters.**
| Hyperparameter | Tomato | Traffic | Glucose | Pandemic |
| --- | --- | --- | --- | --- |
| Training iterations | 500 | 250 | 500 | 260 |
| Batch size | 3000 | 40000 | 100000 | 3860 |
| SGD minibatch size | 128 | 16384 | 1024 | 64 |
| SGD epochs per iteration | 8 | 5 | 4 | 5 |
| Optimizer | Adam | Adam | Adam | Adam |
| Learning rate | 1e-3 | 5e-5 | 1e-4 | 0.0003 |
| Gradient clipping | 0.1 | None | 10 | 10 |
| Discount rate ($\gamma$) | 0.99 | 0.99 | 0.99 | 0.99 |
| GAE coefficient ($\lambda$) | 0.98 | 0.97 | 0.98 | 0.95 |
| Entropy coefficient (start) | 0.01 | 0.01 | 0.01 | 0.1 |
| Entropy coefficient (end) | 0.01 | 0.01 | 0.01 | 0.01 |
| Entropy schedule horizon | 0 | 0 | 0 | 500000 |
| KL target | 0.001 | 0.02 | 1e-3 | 0.01 |
| Value function loss clipping | 10 | 10,000 | 100 | 20 |
| Value function loss coefficient | 0.1 | 0.5 | 0.0001 | 0.5 |
| Share value function layers | F | T | T | T |

**Table 3: ORPO-specific hyperparameters.**
| Hyperparameter | Tomato | Traffic | Glucose | Pandemic |
| --- | --- | --- | --- | --- |
| Discriminator reward clipping | 1000 | 10 | 1e10 | 0.1 |
| Regularization coefficient ($\lambda$) | Varied | Varied | Varied | Varied |
| $\sigma_{\tilde{R}}$ | 0.05 | 0.0002 | 0.05 | 0.08 |

##### ORPO details

We found that a couple of tricks were useful to ensure that ORPO remained stable. First, we clip the discriminator term added to the reward functions to a range $[-\delta,\delta]$, since sometimes it can blow up and cause numerical issues. Second, when estimating $\widehat{\chi^{2}}$, we use a trimmed mean (trimmed by 1% in each tail) to reduce the effect that outliers have on the estimate. These are both particularly important for $\chi^{2}$ divergence, where the output of the discriminator is exponentiated for both the reward discriminator term and for estimating $\widehat{\chi^{2}}$.

### D.2 RLHF experiments

We train LLMs via the RLHF implementation used by Coste et al. (), which is based on OpenAssistant and trlX (Havrilla et al., ). To implement $\chi^{2}$ or KL regularization, we directly add a loss term to the PPO loss:

$$ $\chi^{2}$ divergence: $\displaystyle\lambda\left(\frac{\pi_{\theta}(a\mid s)}{{\pi_{\text{base}}}(a\mid s)}+\frac{{\pi_{\text{base}}}(a\mid s)}{\pi_{\theta}(a\mid s)}-2\right)$ KL divergence: $\displaystyle\lambda\left(\log\frac{\pi_{\theta}(a\mid s)}{{\pi_{\text{base}}}(a\mid s)}+\frac{{\pi_{\text{base}}}(a\mid s)}{\pi_{\theta}(a\mid s)}-1\right)$ $$

where $s$ is the prompt and $a$ is the sampled response. Intuitively, both loss terms have a unique minimum when $\pi_{\theta}(a\mid s)={\pi_{\text{base}}}(a\mid s)$ and in expectation are equivalent to the correct divergence. Schulman () suggests that these are particularly low-variance estimates, and we find that they work well in practice.

## Appendix E Additional Experiments and Results

In this appendix, we the full results from our main experiments as well as ablations of ORPO.

### E.1 Results for all regularization coefficients

Here, we present the results of training with AD and OM regularization using $\chi^{2}$ and KL divergence across all regularization coefficients. Each table shows the median and the standard deviation of the true rewards achieved by the learned policy across 5 random seeds.

**Table 4: All traffic control results ($\times 10^{3}$).**
| Coefficient | AD | State OM | State-action OM |
| --- | --- | --- | --- |
| 0.000002 | -60.79 $\pm$ 8.40 | -61.04 $\pm$ 2.20 | -59.11 $\pm$ 4.53 |
| 0.000004 | -62.01 $\pm$ 8.58 | -61.85 $\pm$ 2.22 | -58.21 $\pm$ 5.52 |
| 0.00001 | -60.43 $\pm$ 6.61 | -59.82 $\pm$ 2.85 | -54.35 $\pm$ 2.81 |
| 0.00002 | -62.01 $\pm$ 4.54 | -56.53 $\pm$ 10.66 | -1.35 $\pm$ 30.28 |
| 0.00004 | -50.84 $\pm$ 6.28 | -42.24 $\pm$ 22.01 | -1.15 $\pm$ 0.06 |
| 0.0001 | -1.29 $\pm$ 0.12 | -2.18 $\pm$ 0.42 | -1.38 $\pm$ 0.29 |
| 0.0002 | -1.70 $\pm$ 0.12 | -2.46 $\pm$ 0.65 | -1.85 $\pm$ 0.29 |

**Table 5: All pandemic mitigation results.**
| Coefficient | AD | State OM | State-action OM |
| --- | --- | --- | --- |
| 0.0008 | -17.60 $\pm$ 1.78 | -12.59 $\pm$ 1.63 | -27.16 $\pm$ 12.83 |
| 0.0016 | -36.16 $\pm$ 3.25 | -11.25 $\pm$ 10.08 | -11.17 $\pm$ 0.19 |
| 0.004 | -34.16 $\pm$ 12.90 | -10.88 $\pm$ 0.91 | -11.65 $\pm$ 2.12 |
| 0.008 | -12.45 $\pm$ 0.25 | -10.99 $\pm$ 3.13 | -12.02 $\pm$ 0.18 |
| 0.016 | -12.31 $\pm$ 0.08 | -10.68 $\pm$ 0.17 | -12.18 $\pm$ 0.46 |
| 0.04 | -12.29 $\pm$ 0.05 | -10.78 $\pm$ 0.12 | -12.34 $\pm$ 0.10 |
| 0.08 | -12.39 $\pm$ 0.04 | -10.73 $\pm$ 0.84 | -12.20 $\pm$ 0.11 |

**Table 6: All glucose monitoring results ($\times 10^{3}$).**
| Coefficient | AD | State OM | State-action OM |
| --- | --- | --- | --- |
| 0.0005 | -580.7 $\pm$ 73.8 | -484.2 $\pm$ 56.6 | -164.0 $\pm$ 3.5 |
| 0.001 | -94.2 $\pm$ 14.0 | -263.9 $\pm$ 19.2 | -127.0 $\pm$ 5.9 |
| 0.0025 | -97.1 $\pm$ 8.2 | -146.4 $\pm$ 18.1 | -93.3 $\pm$ 9.3 |
| 0.005 | -76.6 $\pm$ 10.5 | -109.7 $\pm$ 10.3 | -55.2 $\pm$ 0.7 |
| 0.01 | -84.7 $\pm$ 8.5 | -72.8 $\pm$ 8.8 | -47.5 $\pm$ 0.6 |
| 0.025 | -85.6 $\pm$ 7.9 | -54.3 $\pm$ 1.3 | -50.9 $\pm$ 2.3 |
| 0.05 | -74.8 $\pm$ 13.1 | -57.1 $\pm$ 3.5 | -113.3 $\pm$ 32.8 |

**Table 7: All results for RLHF.**
| Coefficient | AD $\chi^{2}$ | AD KL |
| --- | --- | --- |
| 0.00008 | 9.20 $\pm$ 0.68 | 8.80 $\pm$ 2.24 |
| 0.00025 | 14.05 $\pm$ 3.27 | 9.48 $\pm$ 1.02 |
| 0.0008 | 16.94 $\pm$ 0.07 | 8.84 $\pm$ 0.42 |
| 0.0025 | 16.84 $\pm$ 0.08 | 14.22 $\pm$ 2.81 |
| 0.008 | 16.71 $\pm$ 0.11 | 12.73 $\pm$ 2.75 |
| 0.025 | 16.59 $\pm$ 0.11 | 16.81 $\pm$ 0.27 |
| 0.08 | 16.43 $\pm$ 0.07 | 16.52 $\pm$ 0.08 |
| 0.25 | 16.22 $\pm$ 0.10 | 16.33 $\pm$ 0.04 |
| 0.8 | 16.13 $\pm$ 0.10 | 16.25 $\pm$ 0.13 |

**Table 8: All results for the tomato-watering gridworld.**
| Coefficient | AD | State OM | State-action OM |
| --- | --- | --- | --- |
| 0.0005 | 2.65 $\pm$ 0.67 | 0.64 $\pm$ 0.19 | 0.53 $\pm$ 0.29 |
| 0.001 | 3.87 $\pm$ 0.25 | 0.65 $\pm$ 4.63 | 0.60 $\pm$ 0.16 |
| 0.0025 | 6.24 $\pm$ 0.10 | 9.06 $\pm$ 0.17 | 9.04 $\pm$ 4.71 |
| 0.005 | 6.17 $\pm$ 0.03 | 9.06 $\pm$ 0.11 | 9.16 $\pm$ 0.09 |
| 0.01 | 6.16 $\pm$ 0.04 | 9.07 $\pm$ 0.07 | 9.17 $\pm$ 0.12 |
| 0.025 | 6.19 $\pm$ 0.04 | 8.64 $\pm$ 0.12 | 8.61 $\pm$ 0.18 |
| 0.05 | 6.14 $\pm$ 0.00 | 7.95 $\pm$ 0.05 | 7.99 $\pm$ 0.22 |

### E.2 Ablations

Here, we present the results of two ablations of ORPO. For each of the non-RLHF environments, we fix the optimal coefficient for state-action OM $\chi^{2}$ regularization and modify other hyperparameters. We do not ablate the RLHF experiments because RLHF is a contextual bandit (Appendix A.3) and so it isn’t actually necessary to run ORPO for RLHF. The results are shown in Table 9 below.

Order of training policy and discriminator networks.
First, we experiment with modifying ORPO to train the discriminator *after* the policy. In Algorithm 1, the discriminator $\hat{d}_{\phi}$ is optimized, then the rewards are updated with the discriminator outputs, and then the policy is trained with the updated rewards. An alternative is to wait to train the discriminator until after the policy has been updated, i.e., put Line 4 after Line 7. We experimented with this and found that in in most environments there is not too much difference between the two orders, although training the discriminator first gives slightly better results. However, in the pandemic mitigation environment, we found that it training the discriminator second gave results with much higher variance. This suggests that it is best to train the discriminator before augmenting the rewards to train the policy.

Discriminator reward clipping.
Second, we experiment with modifying the discriminator reward clipping parameter $\delta$ of ORPO. We found that removing the clipping parameter entirely led to NaN errors and training could not complete, so we do not report those results. However, to test sensitivity to this parameter, we tried training with a clipping parameter $10\times$ larger and $10\times$ smaller in each environment. We found that the results did not vary by much across different clipping parameters. This suggests that ORPO is relatively robust to the hyperparameter, so it does not need to be tuned precisely.

**Table 9: Results of our ablations of ORPO. We report the median and standard deviation of the true reward across five random seeds for the four non-RLHF environments. The top row shows the results of state-action occupancy measure regularization with the optimal coefficient for $\chi^{2}$ divergence. The other rows show ablations with the same coefficient. We find that ORPO is mostly robust to different hyperparameters but that it is probably best to train the discriminator network before the policy network.**
|  | Environment |  |  |  |
| --- | --- | --- | --- | --- |
|  | Traffic | Pandemic | Glucose | AI safety |
| Method | control | mitigation | monitoring | gridworld |
|  | $(\times 10^{3})$ |  | $(\times 10^{3})$ |  |
| Default parameters | -1.15 $\pm$ 0.06 | -12.18 $\pm$ 0.46 | -47.5 $\pm$ 0.6 | 9.17 $\pm$ 0.12 |
| Train policy before discriminator | -1.24 $\pm$ 0.07 | -12.23 $\pm$ 2.71 | -48.2 $\pm$ 0.7 | 8.97 $\pm$ 0.13 |
| Discriminator reward clipping $\times 0.1$ | -1.18 $\pm$ 0.14 | -12.22 $\pm$ 0.21 | -47.7 $\pm$ 0.8 | 9.21 $\pm$ 0.18 |
| Discriminator reward clipping $\times 10$ | -1.24 $\pm$ 0.03 | -12.11 $\pm$ 0.29 | -47.9 $\pm$ 0.7 | 9.20 $\pm$ 0.14 |

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 4
questions: ## Typos
there is an issue with the text in Definition 4.2: I think your $\tilde{\mathcal{R}}$ ate something after "policy" and another thing before "lower".

## Suggestions
* line 24: "superior" feels a bit strong
* line 164: the SEIR and glucose environments diagrams are too small to read/understand. just a comic of the environment would be better. the gridworld is on the edge of too small but ok
* line 217: I would avoid using the word "misaligned", since that has all kinds of connotations. Consider just "different", "not identical to", etc. This is a nit: if you feel strongly about "misaligned", I think it's probably ok.
* line 312: the SEIR model has been around for a long time (a quick search found stuff from 2013 but I think it's been around for significantly longer); please try to find an original reference OR change the sentence to talk specifically about that one from 2020 which you're using.
* 369: RHS = right hand side? some readers might not know this; please define on first use
* 402: "widely used to prevent reward hacking in RLHF" well, I don't think this is quite true. KL divergence, to my understanding, is primarily used to avoid instabilities in the training procedure (important in when using PPO in general, not just RLHF!). Now, it seems likely that there is a side-effect of this regularization which helps prevent hacking, but that's not the main reason (in particular: if you remove KL divergence, you don't get hacking, you just get a failure to train). Please update this sentence! (or please explain if I'm missing something).

## Questions
* copying from above, my biggest concern is this one: even if a natural base policy does exist, it's not obvious that the state-occupancy distribution induced by it is the one we actually care about: the optimal policy wrt the true reward might be quite different from what we think of as the "natural" base policy. The issue I see here is that if the "actually really good" (much better than "reasonable base") policies have a way different state-occupancy distribution than the base policy, then the "check" that the proxy reward is "similar enough" to the true reward over the "reasonable" distribution just doesn't matter: what matters is what happens near optimality. Am I missing something here?
* why do you only consider the *optimal* policy wrt the proxy in Definition 4.2? This seems too restrictive (and you apply the regularization at all points during training, so unless I'm missing something you don't rely on this optimality?)
* how do you expect your technique to scale?

## Closing note
I gave the paper a 5 because I think it's important that some things be changed before it's ready for publication. But I will be happy to increase the score to a 6 if my minor points are met, and possibly more if my greater fears about the utility of the technique in capturing reward hacking in different but relevant parts of state-action occupancy space are allayed. Indeed, it's possible that just acknowledging and writing about this issue would be enough: so long as it is clearly framed what this paper is and is not doing and to what extent it "covers all the bases" wrt reward hacking, I think it's worth publishing. It's really nicely presented and a nice idea with experiments to back it.
soundness: 4
strengths: The paper is extremely well-written.

The paper builds on past literature in a coherent and helpful way.

The paper shows experiments on interesting environments.
summary: The authors propose a new definition for reward hacking, based on the idea that we should only care about hacking of proxy reward functions which are correlated (in terms of state-action occupancy) with the true reward function under some "reasonable" base policy.

They then show theoretically and empirically that regularizing reward using an occupancy measure term as opposed to the standard KL divergence leads to less reward hacking across several simple environments.
weaknesses: While the new definition can be defended as an improvement over past definitions, it still leaves some things to be desired. Specifically:

* it requires a "natural base policy", which might not exist in some real-world relevant settings
* even if a natural base policy does exist, it's not obvious that the state-occupancy distribution induced by it is the one we actually care about: the optimal policy wrt the true reward might be quite different from what we think of as the "natural" base policy.
* the definition only considers hacking to have occurred when looking at the *optimal* policy with respect to the proxy; in practice, we never find the optimal policy, yet hacking occurs nonetheless with intermediate sub-optimal (or locally optimal) policies
* the paper doesn't say how to know whether, and what to do if, the (eg, learned) proxy reward function is not r-correlated with the true one.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: ## Questions about underlying assumptions
1. Does the paper assume that both the base policy $\pi_{base}$ and the optimized policy $\pi$ are able to acess the same observation states?
2. **Handling of Unobserved Variables**: In the traffic domain, several potentially impactful variables may be unobserved, as noted in [1][2]. For example, factors like the driver’s expertise level [1] or weather conditions [1] are often missing in driving datasets. How does your framework address scenarios where such unobserved variables exist?
3. During the training phase, does the agent have access to the ground truth reward value?
4. What are major assumptions behind your framework?

## Questions about metrics
1. **Unknown True Reward**: In many practical applications, such as driving, it is challenging to properly define a true reward function [1][2]. In RLHF, the system often relies on preference pairs rather than a true reward. When the true reward is unavailable, what alternative metrics would you recommend for evaluation?
2. **Alternative Metrics for RLHF**: For RLHF, additional metrics might be helpful, such as the winning rate (preferred vs. rejected) or a safety score. How could these be considered in your framework?

## Other concerns
1. **Similarity to Prior Work**: This paper is similar to [3], particularly in the theoretical sections and some experiments. It would be beneficial for the authors to clarify the distinctions between their proposed approach and that of [3].
2. **$\chi^{2}$ vs. KL Divergence**: Does $\chi^{2}$ divergence consistently outperform KL divergence in all scenarios? Are there cases where KL divergence yields better results? If so, what trade-offs are involved? If not, could the authors outline situations where KL divergence might be more suitable?

----------
[1] Ruan, Kangrui, and Xuan Di. "Learning human driving behaviors with sequential causal imitation learning." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. No. 4. 2022.

[2] De Haan, Pim, Dinesh Jayaraman, and Sergey Levine. "Causal confusion in imitation learning." Advances in neural information processing systems 32 (2019).

[3] Laidlaw, Cassidy, Shivam Singhal, and Anca Dragan. "Preventing Reward Hacking with Occupancy Measure Regularization." ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems.
soundness: 3
strengths: The concept presented, while not very complex, demonstrates impressive effectiveness.The authors support their claims with a solid mix of theoretical proofs and comprehensive experiments across various domains. Specific comparison between $\chi^{2}$ vs. KL divergence has been provided. Intuitive examples are also provided. Additionally, intuition about "optimizing the occupancy measure (OM) divergence between policies" might be helpful to the whole community. Overall, this is an excellent paper with good results and insightful contributions.
summary: The paper addresses the issue of reward hacking in RL, where policies optimized with proxy rewards. To address this, the authors propose a novel, formal definition of reward hacking, which is based on the correlation between proxy and true rewards observed by a base policy. The paper also propose a regularization method based on $\chi^{2}$ divergence between the policies' occupancy measures.
weaknesses: However, my primary concern lies in the assumptions underpinning their framework and the choice of metrics. Please see the specific questions below for further details.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: - The definition of correlated rewards allows for the correlated reward to deviate arbitrarily in states with zero occupancy by the base policy.  Doesn't this allow for policies that can be dangerous to still be correlated with the true reward? For example, if the base policy avoids a state with a large negative true reward but the proxy reward function assigns a very high positive reward for that state. Wouldn't this reward incentivize policies that try to visit this state and get a large negative true reward? In addition, wouldn't these kinds of proxy create policies that break the assumption in Theorem 5.1?
soundness: 4
strengths: - The paper is very well written with an astounding flow of information. 
- The authors provide a lot of great intuitive examples for the arguments they make.
- The authors provide a wide range of experiments that show both the soundness and robustness of the claims.
summary: The authors provide a new definition of reward hacking based on a new definition of correlated policies. Reward hacking is a phenomenon that occurs when optimizing a proxy reward for a true reward function yields worse performance over the true reward. The definition of correlated captures proxy reward functions that match the true reward function in visited state action pairs by a base policy. In addition, the authors propose occupancy measure regularization with respect to a base policy as a method to overcome reward hacking. They show theoretically and empirically that the proposed method mitigates reward hacking.
weaknesses: - In Theorem 5.1, there is the assumption that the state-occupancy measure of the policy is absolutely continuous with respect to the base policy. However, during training, this is not something that can be guaranteed. I suggest adding an extra discussion on how regularization also pushes toward the learned policy satisfying this assumption.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: In the introduction, you claim that the formalization of a "good proxy"
has been elusive. Hasn't this issue been addressed by the EPIC distance
\[1\]? What does EPIC leave to be desired?

Is there a mistake in the Table 1 cell for traffic control with state
occupancy KL? Particularly, the standard deviation of $22.6$ is
alarming.

Theorem 5.1 is nice in the sense that it tells us that occupancy measure
regularization (particularly w.r.t $\chi^2$) can prevent reward hacking
even when the proxy is correlated with the true reward. However, there
is no claim about whether or not reward hacking can (at least in
principle) be prevented with action distribution regularization. So,
naturally, I'm curious whether or not a similar result can be achieved
with action distribution regularization. The empirical results suggest
not, but perhaps we just haven't found the right way to do it yet. Do
you think it's impossible?

In the experiments, RLHF is treated as a contextual bandit problem, in
which case there is no difference between occupancy measure
regularization and action distribution regularization (because state
occupancy is independent of actions). However, especially given the
autoregressive nature of many existing LLMs, RLHF can also be viewed as
a sequential problem where the policy sequentially chooses individual
tokens, and a reward is presented after the last token is chosen. How
would OM under this formulation of RLHF compare to AD in the "atomic"
contextual bandit case?

Theorem 5.1 looks nice because intuitovely, it tells us that by
controlling the $\chi^2$ between the learned policy and the base policy,
we can prevent reward hacking. But(!), it is actually not entirely clear
to me if this lower bound admits the possibility that we can actually
*improve* upon the base policy using this regularization. Clearly, as
the $\chi^2$ divergence tends to $0$, the lower bound tends to $0$. But
can the lower bound ever be greater than $0$?

## References

1.  Gleave, Dennis & Legg et al. (2021) Quantifying Differences in
    Reward Functions, arXiv.
2.  Fujimoto, Meger & Precup (2019) Off-Policy Deep Reinforcement
    Learning Without Exploration, ICML.
soundness: 3
strengths: This paper studies a very hot topic in the field, and in my opinion,
provides some very interesting insight. Most notably, I really enjoyed
the following entangled results in the paper:

1.  Even with a proxy reward that is highly correlated with the ground
    truth reward, reward hacking is possible;
2.  But under such correlated rewards, it is (probably) possible to
    prevent reward hacking by regularizing the fine-tuned policie's
    occupancy measure (wrt $\chi^2$).

It is also nice that their proposed framework fits nicely with that of
IRL, allowing the authors to leverage performant IRL algorithms for RL
fine-tuning.
summary: This paper critically formalizes a notion of exploitability in RL-based
fine-tuning when the reward function is difficult or impossible to
specify exactly, and is instead replaced by a "proxy" reward. It argues
that reward hacking is characterized by the proxy reward admitting
optimal policies that are in fact *not* aligned with our desired
behavior, and their characterization is validated with several intuitive
case studies. Moreover, under this characterization, the authors propose
a RL fine-tuning algorithm that prevents reward hacking by regularizing
the *occupancy measure* of the fine-tuned policy towards a base policy
with respect to $\chi^2$ divergence, under weaker assumptions on the
proxy reward than existing results. Moreover, the paper introduces a
practical implementation using techniques from inverse RL (particularly
GAIL), and demonstrate that their method performs better than baselines
across several fine-tuning tasks.
weaknesses: My main concerns with the paper regard lack of clarity in several
respects. I will outline these below, and further examples are given in
"Minor Issues" and "Questions".

1.  I find the narrative of the paper to be a little misguided. The
    paper goes on to try to essentially axiomatize reward hacking,
    leading to many fairly philosophical or imprecise heuristic
    discussions (e.g. "it's not reward hacking if the proxy is just
    bad!"). I think this is, in a sense, actually overcomplicating
    things. Under the hood, I think there is a very clear and
    uncontroversial message here:
    1.  When fine-tuning, we want to learn a policy that is better than
        some base policy at maximizing some ground truth reward (this is
        hard to argue with);
    2.  Proxy rewards that are highly correlated to the ground truth
        reward are still "hackable", in the sense that policy
        optimization with such proxies prohibits improvement over the
        base policy (the authors show this);
    3.  But hackable correlated proxies are not useless—it is possible
        to successfully fine-tune a base policy with such a proxy using
        $\chi^2$ regularization (the authors basically show this, see
        Questions).
2.  While interesting, I think some of the theoretical results are
    actually fairly weak (see Questions). For instance, while the
    theoretical results suggest regularization of the fine-tuned
    policy's occupancy measure, they do not actually suggest that this
    should be an improvement over action-distribution regularization.
    Moreover, it is not actually clear that the lower bound of Theorem
    5.1, which is the main theorem of the paper, can actually be
    positive (in other words, Theorem 5.1 does not explicitly show that
    fine-tuning with occupancy measure regularization can actually lead
    to improved alignment). I suspect it is possible for the authors to
    correct these, and that would strengthen the paper.
3.  Overall, there are several issues with writing, such as discussion
    of related work, redundancy in intro/background sections, and some
    terminology, which I outlined below.

My score is mostly a reflection of weakness 2 here. I trust that the authors can fix weaknesses 1 and 3 fairly easily. Should the authors address weakness 2 (as well as the questions below), I would be happy to increase my score.

## Minor Issues

Stiennon et al. (2020) is definitely not the first example of using KL
regularization to an "offline policy". This is very commonly done in all
of maximum entropy / entropy-regularized reinforcement learning, dating
back at least fifteen years before this reference.

With regard to prior work in offline RL that enforce that the learned
policy does not veer too far away from the data distribution, the work
of \[2\] is the most appropriate reference (see list of references in
Questions section below).

The paper starts off much too slow. The first four pages are fairly
redundant, the isssue of formalizing the notion of reward hacking is
addressed several times, and only very heuristic high-level ideas are
suggested for addressing the issue. I believe this pages can be
condensed a lot, which would make the paper easier to read.

The section about the desiderata for a definition of reward hacking can
be clarified a lot. In particular, it is not clear to me exactly what
the desiderata are, they should be stated very explicitly (especially
since the paper later refers to them in some order, e.g. "the second
desideratum").

At the bottom of page 5, $\sigma_{\tilde{R}}$ and $\sigma_R$ are really
variances, not standard deviations (so they should probably be written
as $\sigma_{\bullet}^2$).

In definition 4.2, there is a missing word — "*reward hacking* occurs
when an optimal policy $\tilde{R}$ **lower** true reward…".

In equation (8), you use $\phi$ on both sides of the equation (you
should use e.g. $\phi'$ as the variable being minimized).

While definition 4.2 is interesting, I don't think "reward hacking" is
the right name for the phenomenon being described in this definition.
Particularly, this definition presents a property on a proxy reward, but
not on the process of fine-tuning a policy with this proxy. Indeed, one
of the contributions presented later in the work is an algorithm that
prevents reward hacking under proxies satisfying definition 4.2! Rather,
this definition presents conditions under which reward hacking *may*
occur. As such, I believe "reward exploitability" is an example of a
more appropriate name: standard policy optimizations *can* exploit/hack
the reward function if you're not careful.

On line 313, "polulation" should be "population".

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: Please see my comments and questions *Weaknesses*. I'd appreciate if the authors could respond to these comments and questions.
soundness: 3
strengths: The paper is well-written, enjoyable to read, and relatively easy to follow. It provides an interesting perspective on reward hacking through the lens of regularized RL objectives. The approach is based on a new definition of reward hacking, which assumes access to a base policy-a reference used to determine whether a proxy reward is hackable by a learning agent. This definition is intuitive, although I have some questions and comments about its expressiveness.

I find the connection to the KL-regularized RLHF objective, and the corresponding result (Theorem 3.1), particularly interesting. To the best of my knowledge, this result is novel. The experimental results are extensive and, in my opinion, demonstrate the effectiveness of the proposed approach.
summary: The paper studies reward hacking, including in the context of reinforcement learning from human feedback. It proposes a new definition of what constitutes reward hacking and shows that a specific form of regularization can help with preventing reward hacking. In the context of RLHF, the results suggest a different regularization scheme than a standard KL penalty. The authors further argue that this type of regularization better mitigates reward hacking in practice.
weaknesses: Regularization to a base policy in RLHF for language models is somewhat natural, as we aim to preserve the language generation capabilities of our policy. However, as shown in Table 1, it can also lead to considerable performance loss in some environments (e.g., Pandemic Mitigation), where the base policy does not perform well. I wonder how preference-based RL (without regularization) would fare in such environments. It seems that the proposed approach is somewhat constrained, as it requires humans to provide a reliable reference point (base policy), which may be challenging in complex environments.
 
While I am generally positive about this work, I think I didn’t fully understand how it compares to some other well-known approaches to mitigating similar problems (e.g., techniques based on inverse reward design or risk-averse RL). I understand that these approaches do not necessarily address the same problem, but the experimental setup does not seem to include baselines that would illustrate the importance of Definition 4.2. For example, since $\pi_{base}$ is given, why isn’t IRL an adequate baseline?
 
The related work section covers most of the important references, but a more comprehensive literature review on RLHF could be beneficial. For example, the paper mentions:

> Our regularized objective in (5) differs in two main ways from the KL regularization used in RLHF. First, our results suggest optimizing the occupancy measure (OM) divergence between policies, whereas RLHF uses the action distribution (AD) divergence. 

however, it does not explain how this connects to recent results that propose OM-based regularizers. For example, it would be useful if the authors could comment on the MDP formulation of RLHF from

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

which in Section 7.2 introduces the KL-regularized RHF objective based on occupancy measures. 

One potential drawback of the proposed approach is that it is based on estimating occupancy measures, which may be challenging to estimate for high-dimensional state spaces. It would be helpful if the authors could comment on the level of difficulty of the current environments for an RL agent. Furthermore, how accurate are the occupancy measures obtained in the experiments? It seems that it should be possible to estimate this for some of the environments.

Some minor details in the experimental setup were not clear to me. The appendix states:

> We primarily tuned the hyperparameters listed below in order to ensure that the proxy reward would be properly optimized and reward hacking would occur without regularization.

Does this mean that we typically don't observe reward hacking in these environments?

Meta-review: The reviewers agreed that the paper investigates an important problem of reward hacking, proposes a novel definition, and then develops a new regularization that can effectively prevent reward hacking. However, the reviewers also raised several concerns and questions in their initial reviews. We want to thank the authors for their responses and active engagement during the discussion phase. The reviewers appreciated the responses, which helped in answering their key questions. The reviewers have an overall positive assessment of the paper, and there is a consensus for acceptance. The reviewers have provided detailed feedback, and we strongly encourage the authors to incorporate this feedback when preparing the final version of the paper.

Decision: Accept (Spotlight)