# Goodhart in RL with KL: Appendix

This is the appendix to the [previous post](https://www.lesswrong.com/s/6rhjdbnEXoek4YiH7/p/pEZoTSCxHY3mfPbHu) on Goodhart’s Law and KL regularization, containing all of our proofs.

Theorem about distributions
---------------------------

***Theorem 1:*** Given any heavy-tailed reference distribution $Q$ over $\mathbb R$ with mean $\mu_Q$, and any $M, \epsilon > 0$, there is a distribution $P$ with mean $\mu_P>M$ and $D_{KL}(P \| Q) < \epsilon$.

*Proof:* WLOG let $\mu_Q = 0$. We construct a sequence of distributions $\{P_t\}$ such that $\lim_{t \to \infty} \mathbb E_{P_t}[X] \ge c$ for any constant $c$, and $\lim_{t \to \infty} D_{KL}(P_t \| Q) = 0$. We define $P_t$ for any $t > c$ thusly. Writing $F_{P_t}(x)$ for the CDF $Pr_{X \sim P_t}(X \le x)$ and $\bar F_{P_t}(x)$ for $1 - F_{P_t}(x)$, we let

$$
\bar F_{P_t}(x) = \begin{cases} 1 - \frac{1 - c/t}{F_Q(t)}F_Q(x) & x \le t \ \frac{c/t}{\bar F_Q(t)}\bar F_Q(x) & x > t \end{cases}
$$

Intuitively, we rescale the part of the distribution to the right of $t$ evenly to have total probability $c/t$, which is less than 1 because $t > c$.

We must check that $\lim_{t \to \infty} \mathbb E_{P_t}[X] = c$. We can write

$$
\begin{align*} \mathbb E_{P_t}[X] = F_{P_t}(t) \mathbb E_{P_t}[X | X \le t] &+ \bar F_{P_t}(t) \mathbb E_{P_t}[X | X > t]    \ = F_{P_t}(t) \mathbb E_Q[X | X \le t] &+ \bar F_{P_t}(t) \mathbb E_Q[X | X > t]    \ = F_Q(t) \mathbb E_Q[X | X \le t] &+ \bar F_Q(t) \mathbb E_Q[X | X > t] + \ (F_{P_t}(t) - F_Q(t))&E_Q[X | X \le t] + (\bar F_{P_t}(t) - \bar F_Q(t)) E_Q[X | X > t]     \ = \mathbb E_Q[X] + (\bar F_{P_t}(t) - &\bar F_Q(t))(E_Q[X | X > t] - E_Q[X | X \le t])    \end{align*}
$$

We know that $\mathbb E_Q[X | X > t] > t$ because it is an integral of values strictly greater than t. Because $\mathbb E_Q[X] = 0$ is a weighted average of $\mathbb E_Q[X|X>t]$ and $E_Q[X | X \le t]$, and $\mathbb E_Q[X|X>t] > 0$, we know $E_Q[X | X \le t] < 0$. So $E_Q[X | X > t] - E_Q[X | X \le t] > t.$ We also know that for sufficiently large $t$, $(F_{P_t}(t) - F_Q(t)) > 0$. Intuitively, starting from $Q$, which has mean 0, $P_t$ moves a probability mass approaching $\frac c t$ from mean <0 to mean >t.

Now we can say

$$
\lim_{t \to \infty} \mathbb E_{P_t}[X] > \lim_{t \to \infty}\left[ \mathbb E_Q[X] + (\bar F_{P_t}(t) - \bar F_Q(t))(t-0) \right] \ = \lim_{t \to \infty}\left( \frac c t - \bar F_Q(t) \right) t = \lim_{t \to \infty}c - t \bar F_Q(t)
$$

Because $Q$ has a finite mean, $\lim_{t \to \infty} t \bar F_Q(t) = 0$, and so $\mathbb \lim_{t \to \infty} \mathbb E_{P_t}[X] \ge c$.

Now we check that $\lim_{t \to \infty} D_{KL}(P_t \| Q) = 0$:

$$
\begin{align*}D_{KL}({P_t} \| Q) &= \int_{\mathbb R} \log \frac{{P_t}(dx)}{Q(dx)} \,{P_t}(dx)\&= \int_{x \le t} \log \frac{{P_t}(dx)}{Q(dx)} \,{P_t}(dx) + \int_{x > t} \log \frac{{P_t}(dx)}{Q(dx)} \,{P_t}(dx)\ &= F_{P_t}(t) \log \frac{F_{P_t}(t)}{F_Q(t)} + \bar F_{P_t}(t) \log \frac{\bar F_{P_t}(t)}{\bar F_Q(t)} \text{\quad since both ratios are constant}\ &= F_{P_t}(t) \log \frac{1 - c/t}{F_Q(t)} + \bar F_{P_t}(t) \log \frac{\bar F_{P_t}(t)}{\bar F_Q(t)}\end{align*}
$$

Since both $1-c/t$ and $F_Q(t)$ go to $1$ as $t \to \infty$, the left term goes to $0$, and so

$$
\begin{align*}\lim_{t \to \infty} D_{KL}(P_t \| Q)&\le 0 + \lim_{t \to \infty}\bar F_{P_t}(t) \log \frac{\bar F_{P_t}(t) }{\bar F_Q(t)} \ & = \lim_{t \to \infty}\frac c t \log \frac c {t \bar F_Q(t)} \le \lim_{t \to \infty} \frac c t \log \frac 1 {\bar F_Q(t)} \&= \lim_{t \to \infty}-\frac c t \log \bar F_Q(t) \quad \text{ since t>c} \end{align*}
$$

$Q$ is heavy tailed, so by definition $\lim _{t \rightarrow \infty} e^{a t} \bar{F}_Q(t)=\infty \quad \text { for all } a>0$. This implies that for every $a > 0$ there is a sufficiently large $t_{c}$ so that for all $t > t_c$, $\bar F_Q(x) > e^{-at}$, which means that $\log \bar F_Q(t) > -a t$.

Therefore for every $a > 0$, $\lim_{t \to \infty} D_{KL}(P_t \| Q) \le \lim_{t \to \infty} -c/t \log \bar F_Q(t) < \lim_{t \to \infty} -\frac {-act} t = ac$, which since KL divergence is nonnegative means that $\lim_{t \to \infty} D_{KL}(P_t \| Q) = 0$ as desired. $\blacksquare$

Theorem about deterministic Markovian-return MDPs
-------------------------------------------------

***Definition:*** A deterministic-transition MDP with Markovian returns (DMRMDP) is an MDP $(\mathcal S, \mathcal A, P, R)$ such that:

*   The transition function $P: \mathcal{S} \times \mathcal{A} \to \mathcal{S}$ is deterministic, i.e., for each state $s \in \mathcal{S}$ and action $a \in \mathcal{A}$, there exists a unique state $s' \in \mathcal{S}$ such that $P(s' | s, a) = 1$.
*   There is a set of sink states $E \subseteq \mathcal S$ that terminate a trajectory, which is disjoint with the set of start states.
*   Returns are Markovian; that is, for any two trajectories $\tau=(s_1, a_1, \dots, s_n), \tau'=(s_1', a_1', \dots, s_n'),$ if $s_{n}= s'_{n}$, then $\tau$ and $\tau'$ have identical return distributions. Equivalently, for the trajectory random variable $T=(S_1, A_1, \dots)$ distributed according to any policy, with return $G$, $G \bot\!\!\!\!\! \bot (S_{<i}, A_{<i})  | S_i$ for any $i \ge 1$.

Note: Sampling from a language model and applying RLHF is well-modeled as a DMRMDP, since the state is a sequence of tokens (actions) which deterministically results from the last token and returns depend only on the final state.

***Theorem 2:*** Let $W = (\mathcal S, \mathcal A, P, R)$ be a deterministic-transition MDP with Markovian returns. Given $W$ we define the function that takes policies to trajectories $Tr: (S \to \Delta A) \to \Delta(S \times A)^*$, and the average return function $g: (S \times A)^* \to \mathbb R$ which induces a function $G: \Delta(S \times A)^* \to \mathbb \Delta \mathbb R$. Let $\pi_0: \mathcal S \to \Delta \mathcal A$ be some reference policy. If $G \circ Tr(\pi_0)$ is heavy-tailed with finite mean $\mu_Q$, then for any $M, \epsilon > 0$, there is a policy $\pi$ with mean return $\mathbb E[U | U \sim G \circ Tr(\pi)] > M$ and $\mathbb E_{s \in T, T\sim Tr(\pi)}[D_{KL}(\pi(s) \| \pi_0(s))] < \epsilon$.

*Proof:* We will exhibit a distribution of trajectories $\rho$ such that $D_{KL}(\rho \| Tr(\pi_0)) < \epsilon$ and $\mathbb E[G(\rho)] > M$, and then construct a policy $\pi$ with $Tr(\pi) = \rho$. Note that this proof applies for continuous action spaces if trajectories are replaced with measurable sets, but this would make it harder to read.

Let $\rho_{\pi_0} = Tr(\pi_0)$. We have a heavy-tailed distribution of return $Q \triangleq G(\rho_{\pi_0})$ over $\mathbb R$, so we can apply Theorem 1. But to define $\rho$, we can construct $P_t$ in the proof of Theorem 1 in a particular way. For any $t>c$, we need a $P_t$ that uniformly upweights values of mean return such that $\bar F_{P_t}(t) = c/t$. We can define $\rho_t$ such that any trajectory $\tau$ is upweighted by a factor depending only on its mean return:

$$
\rho_t(\tau) = \begin{cases} \frac{1 - c/t}{F_Q(t)} \rho_{\pi_0}(\tau) & g(\tau) \le t \ \frac{c/t}{\bar F_Q(t)}\rho_{\pi_0}(\tau) & g(\tau) > t \end{cases}
$$

Then we can let $P_t \triangleq G \circ \rho_t$ and the rest of the proof of Theorem 1 applies. Therefore, applying the theorem, we can let $\rho = \rho_t$ for sufficiently large $t$, and then $\mu_{G \circ \rho} > M$ and $D_{KL}(G \circ \rho, G \circ \rho_{\pi_0}) < \epsilon$. But by the chain rule for KL divergence, $D_{KL}(\rho, \rho_{\pi_0}) = D_{KL}(G\circ \rho, G \circ \rho_{\pi_0}) + \mathbb E_{\gamma \sim G\circ\rho}[D_{KL}(\rho(T) | G(T)=\gamma  \| \rho_{\pi_0}(T) | G(T)=\gamma)]$. Since we constructed $\rho$ so that the probabilities of each $\tau$ conditional on its return being $\gamma$ are equal, the second term is zero, and we also have $D_{KL}(\rho, \rho_{\pi_0}) < \epsilon$.

Finally, since the KL divergence between trajectory distributions is the sum of KL divergence between policies at each action in the trajectory, and each trajectory has at least one action, $\mathbb E_{s \in T, T\sim Tr(\pi)}[D_{KL}(\pi(s) \| \pi_0(s))] \le \mathbb E_{T\sim Tr(\pi)} \sum_{s \in T}[D_{KL}(\pi(s) \| \pi_0(s))] = D_{KL}(\rho \| \rho_{\pi_0}) < \epsilon$ as desired.

To define $\pi$ such that $Tr(\pi) = \rho$, we let $\pi(s, a) = Pr(a_i = a | \tau = (..., s, a_i, ...) \sim \rho)$.

Then the probability that any trajectory $\tau = (s_1, a_1, \dots, a_n)$ is sampled is:

$$
\begin{align} Tr(\pi)(\tau) &= \prod_{i=1}^n \pi(s_i, a_i) \&= \prod_{i=1}^n Pr(a_i=a_i' | \tau' = (..., s, a_i', ...) \sim \rho) \&= \prod_{i=1}^n Pr(a_i = a_i' | \tau' = (s_1', a_1', ..., s, a_i', ...) \sim \rho, s_{<i} = s'{<i}, a{<i} = a'_{<i})\&= \rho(\tau) \end{align}
$$

In (2), returns are Markovian, so all trajectory prefixes ending in state $s$ have the same distribution of returns under any policy. In the construction of $\rho$, all trajectories with the same mean return have equal measure. Therefore, conditioning on earlier states and actions of $\tau$ does not change the measure, so we can write (3). So $Tr(\pi)=\rho$ as desired. $\blacksquare$

Lagrange multipliers to minimize KL divergence
----------------------------------------------

***Theorem 3*****:** If $V$ is light-tailed, $\mathbb E_Q[V]$ is finite, and $d = D_{KL}(P \| Q)$ is bounded, then $\mathbb E_P[V]$ is bounded, and $\mathbb E_P[V] \to 0$ as $d \to 0$.

Using Lagrange multipliers, we find that when KL divergence is minimized, we have $P(V) [\lambda_1 \log \frac{P(V)}{Q(V)} + \lambda_2 - X] = 0$ for some constants $\lambda_1, \lambda_2$, so

$$
\log \frac{P(V)}{Q(V)} = \frac{V - \lambda_2}{\lambda_1}
$$

$$
P(V) = Q(V)\exp\left(\frac{V - \lambda_2}{\lambda_1}\right) = Q(V) e^{V/\lambda_1} e^{-\lambda_2/\lambda_1} = C Q(V) e^{V/\lambda_1}
$$

That is, the new PDF is an [exponential tilting](https://en.wikipedia.org/wiki/Exponential_tilting) of the old PDF. Now what is $\mathbb E_P[V]$? It’s just $\int_{-\infty}^{\infty} C V e^{V/\lambda_1} Q(X) \,dV$. If the distribution of V is heavy-tailed distribution, this is $\infty$; if it is light-tailed, this is some finite value.

When $d = 0$, $P$ and $Q$ are identical and $\mathbb E[V] = 0$. So by a continuity argument, $\mathbb E_P[V] \to 0$ as $d \to 0$. $\blacksquare$

Light tails + independence imply $\mathbb EV \to \infty$
--------------------------------------------------------

***Theorem 4*****:** If $U=X+V$with $X$ and $V$ both light-tailed, and the distribution of U is continuous, and $\pi^*(\beta) \stackrel{\triangle}{=} \arg \max_\pi \mathbb E[U(\pi)] - \beta D_{KL}(\pi, \pi_0)$, then $\lim_{\beta \to 0^+} \mathbb E[V(\pi^*(\beta))] = \infty$.

*Proof:* Fix some $\beta$. Using Lagrange multipliers, we find that for any event $S$, $\Pr_\pi(S) = \Pr_{\pi_0}(S) e^{\lambda U(S)}$. Let $c(\beta)$ be the median value of $U$ under the policy $\pi^*(\beta)$; that is, $Pr(U > c(\beta) | U \sim G \circ Tr(\pi^*(\beta))) = \frac 1 2.$ This exists because $U$ has a continuous distribution. Then:

$$
\begin{align*}E[V | \pi] &= \frac 1 2 E[V | \pi, U < c] + \frac 1 2 E[V | \pi, U \ge c] \ &\ge \frac 1 2 E[V | \pi, U < c] + \frac 1 2 E[V | \pi] \ \lim_{\beta \to 0^+} E[V | \pi] &\ge \lim_{\beta \to 0^+} \frac 1 2 E[V | \pi, U < c] + \lim_{\beta \to 0^+} \frac 1 2 E[V | \pi] \end{align*}
$$

The left term is $c$, while the right term is $\infty$, so the overall limit is $\infty$. $\blacksquare$