Reinforcement Learning from LLM Feedback to Counteract Goal Misgeneralization

Langosco et al. (2023) define goal misgeneralization for a reinforcement learning agent as an instance where a learnt policy $\pi$ achieves a low test reward even though it was achieving a high training reward. The policy appears to be optimizing a different reward than the one originally trained for. More precisely, let $G$ be our intended goal. Consider the case where the agent is deployed in a new environment that presents a distributional shift. Then, $\pi$ is subject to goal misgeneralization if the reward appears to be optimizing a goal $G^{\prime}\neq G$ in the new test environment. We call $G^{\prime}$ the behavioral objective of the agent. In this new environment, the policy continues to act capably yet achieves a low reward. For goal misgeneralization to happen, it must be case that either the training environment is not diverse enough or that that there exists some proxy that correlates with the intended objective during the training, but comes apart once the agent is deployed in the test environment.

Consider a set $D$ of preferences over policy rollouts. Each element of $D$ is triple $(\omega^{1},\omega^{2},\mu)$ where $\omega^{1}$ and $\omega^{2}$ are the two policy roll-outs in a pair, and $\mu$ is a distribution over $\{1,2\}$ indicating which roll-out is preferred. $\mu$ concentrates its mass entirely on the preferred choice or remains uniform if there’s no preference or equal preference for both options. Following Christiano et al. (2017), we train a reward model $\mathcal{R}$ from the set $D$ by minimizing the cross-entropy loss between the predicted probabilities $\hat{P}$ and the preference labels $\mu$:

$$\mathcal{L}(\mathcal{R})=-\sum_{(\omega^{1},\ \omega^{2},\ \mu)\in D}\mu(1)\log\hat{P}(\omega^{1}\succ\omega^{2})+\mu(2)\log\hat{P}(\omega^{2}\succ\omega^{1})$$

where

$$\hat{P}(\omega^{1}\succ\omega^{2})=\frac{\exp\sum_{\omega^{1}}\mathcal{R}(s_{t},s_{t+1})}{\exp\sum_{\omega^{1}}\mathcal{R}(s_{t},s_{t+1})+\exp\sum_{\omega^{2}}\mathcal{R}(s_{t},s_{t+1})}.$$

We consider an RL agent that is prone to goal misgeneralization. This agent is trained for $\tau\cdot T$ timesteps with a specific reward function $R$. While this reward function correctly captures the intended goal during the training, we expect the agent to fail to generalize its goal in out-of-distribution environments because of the existence of a proxy between the intended goal and the agent’s learnt goal. We rely on a Large Language Model (LLM) to identify and correct these failure modes by analyzing the behavior of the RL agent. In particular, we train the agent using a reward model derived from the LLM feedback. The goal of this fine-tuning period is to increase the agent’s generalization capability. The full methodology of this reinforcement learning from LLM feedback is detailed in the next section.

The RL agent is trained for $\tau\cdot T$ timesteps on an environment prone to goal misgeneralization. At $\tau\cdot T$ timesteps we pause the training and sample $N^{I}$ policy roll-outs $\omega^{\pi}$. A roll-out $\omega^{\pi}$ is a sequence of state-action pairs generated by the agent’s policy $\pi$ at time $\tau\cdot T$.

Similarly to how parents observe their children behaviours in different environment and then implement different education methods to correct them, we rely on the assistance of an LLM parent to correct the biased behavior of our RL agent. The objective of LLM supervision is to enable the RL agent to not only accurately generalize its goals but also to maintain competent performance. We input the $N^{I}$ samples into the LLM, prompting it to analyze potential scenarios where our current training setup might cause the agent to inadequately generalize its goal across diverse testing environments. The input to the LLM is structured as follows :

1. 1.

   Preamble A description of the RL problem including the objective of goal generalization

2. 2.

   Request for analysis A question to elicit the LLM into thinking about how the current setup could fail goal generalization

3. 3.

   Explanation of the setup Description of the environment

4. 4.

   Policy rollouts Policy rollout samples

The LLM then suggests modification to the training environment that would mitigate goal misgeneralization. At this point, a straightforward solution to implement the LLM feedback would be to continue the training on the suggested diversified environments directly. Nonetheless, we hypothesize a context where the underlying reward function of these environments cannot be accessed or where the training environment cannot be adjusted. Hence, we aim to correct the agent’s behavior after it has learnt a biased policy preventing it from generalizing its goal. In order to correct this bias, we plan to train our agent using a reward model induced from the LLM preferences. We generate the LLM’s suggested environments and deploy our agent on these environments. We sample $N^{R}$ policy roll-outs. The sole purpose of this training is to construct the LLM-informed reward model.

From the $N^{R}$ policy roll-outs sampled, we form a set comprising all possible pairs. We give the LLM a prompt and two candidate policy roll-outs and and instruct it to rate which policy roll-out is preferred. The input to the LLM for this preference labeling is structured as follows:

1. 1.

   Preamble A description of the RL problem including the objective of goal generalization

2. 2.

   Indication of reward model (optional) An indication that the LLM’s answer will be integrated to a reward model for the RL training

3. 3.

   Explanation of the maze textual representation A key for the textual representation of the maze

4. 4.

   Policy rollouts Policy rollout samples in their textual representation

5. 5.

   Ending An ending string to ensure the LLM answer is in a fixed format (e.g. "Preferred traejctory=")

The LLM preferences are stored in a database $D$ of triples $(\omega^{1},\omega^{2},\mu)$ where $\omega^{1}$ and $\omega^{2}$ are the two policy roll-outs in a pair, and $\mu$ is a distribution over $\{1,2\}$ indicating which roll-out the LLM preferred. If the LLM has a preference, then $\mu$ puts all its mass on that choice. On the contrary, if none is preferred, or if both are equally preferable, then $\mu$ is uniform.

We train a reward model $\mathcal{R}$ from the set $D$ of LLM-preferences. We then resume the training of the RL agent on the training environment. We experiment with combining the LLM-induced reward model with the original reward function. The composite reward function is implemented as:

$$R^{{}^{\prime}}=\lambda\cdot\mathcal{R}+(1-\lambda)\cdot R$$

In this case, implementing our method fully correspond to $\lambda=1$.

We use the OpenAI Procgen environment suite (Cobbe et al., 2019), which has been designed to study RL agents generalization capacities for specific tasks. We choose the Maze game environment where a mouse is trained to navigate through a maze towards a randomly located cheese. A reward function is integrated within the maze game environment and provides a reward of $10$ when the agent reaches the cheese’s location, and a reward of $0$ for any other action.

We follow Langosco et al. (2023) setup to create a setting in the maze game that is prone to goal misgeneralization. In particular, we randomly locate the goal within a region of size $1$ to $10$ in the upper right corner of the maze during the training environment. However, in the test environment, the goal is randomly located in any region within the maze (see figure 1).

We follow a zero-shot protocol in all experiments: the agent does not see the out-of-distribution testing environment during training. Langosco et al. (2023) demonstrated that RL agents trained with the reward function in environments where the goal is consistently placed in the upper right corner learn to navigate there. The learnt objective $G^{\prime}$ is the position (upper right) rather than the intended objective $G$, which is the cheese. However, as the randomization region grows, accurate goal generalization becomes more likely.

We evaluate the Precision, Recall and F1 score for each class on the validation dataset. For each rollout pairs, the LLM preferences are in: $\{(0,1),(0.25,0.75),(0.5,0.5),(0.75,0.25),(1,0)\}$ and represents the strength in preference of rollout 1 or rollout 2. However, due to the very low prevalence of $(0.25,0.75)$ and $(0.75,0.25)$ preferences, each constituting only $0.8\%$, we present our metrics both as an average over all labels and specifically for the predominant classes: ${(0,1),(0.5,0.5),(1,0)}$. in Table 1.

Table 1 shows that macro-average metrics are significantly influenced by the low-prevalence classes ²²2See B.1 for more details on the classes.. The high precision but lower recall for the $(0,1)$ and $(1,0)$ classes indicates that the LLM-reward model is often cautious in preferring one rollout over another, but accurate when it does. This tendency is beneficial, as it helps filter noise from the LLM labeling. Indeed, when the LLM-reward model drives the RL agent’s training, even though the occurrences of high reward are rare, they are informative.

We assess our method on goal randomization regions ranging from size $0$ to size $10$. Our method effectively reduces goal misgeneralization across 7 regions and matches the performance of the existing reward function in three others, as shown in 6.

Our methodology consistently produces an agent that acts capably accross all randomization regions, as it always achieves the maximum mean training reward (10). For randomization regions $\{2,3,4,5,6,7,9\}$ our method effectively increases the agent’s generalization capacity as can be seen from its test reward being greater than with the environment reward function. The success of our method relies on the fact that the LLM-induced reward model captures a penalty for when the agent consistently navigates to the upper right corner despite the cheese being located elsewhere. Our method incentivizes exploration and better environmental cue recognition thus leading to better generalization in out-of-distribution settings.

The most pronounced improvement is observed in randomization region $4$, where the agent’s learning is still flexible, allowing the reward model to more effectively influence its behavior. In contrast, with a random region size of 1, where intended and proxy goals are indistinguishable, the LLM-informed reward model ($\lambda=1$) does not enhance generalization. The disparity in test performance also lessens with larger random regions as the base agent independently learns a generalizable goal.

For training scenarios with $\lambda=0.5$ in randomization regions smaller than 4, our method matches the environmental reward function in terms of generalization. In these cases, the intended and proxy goals are hard to distinguish so the LLM-informed model’s infrequent and mild penalties for biased behaviors are overshadowed by the environmental reward. However, this disparity fades with larger random regions, where the LLM-informed model becomes the primary driver of the agent’s behavior.

One of the key foundations of our method relies on the fact that the LLM parent is able to assess the behavior of the main agent on two dimensions: its ability to fulfill the intended goal, and its capacity to generalize in different environments. Hence the LLM needs to be able to assess whether a particular policy acts toward the intended goal as well as hypothesize on potential aspects it could fail under distributional shifts. In particular, GPT-4 needs to be able to verify whether a particular trajectory could get the agent to the goal in the specific maze environment. In 77% of the times, GPT-4 was correctly able to assess whether a specific trajectory was getting to the goal. ³³3Value obtained by providing 100 examples of mazes and trajectories and asking GPT-4 to verify whether the agent reached the goal. 50 of these examples were trajectories that reached the goal. This contrasts with the 10% of times GPT-4 was able to solve the maze itself. ⁴⁴4Value obtained by providing 50 examples of mazes to GPT-4, and prompting it to solve the maze. All the successes were on mazes that required 3 or less moves. This highlights that, in our case, for LLM supervision, it suffices that the parent LLM be able to assess a specific policy relative to an intended goal, without necessarily being capable of the action itself. This aspect is particularly promising for scalable oversight: LLM can supervise other AI agents towards achieving their tasks and can help them generalize their goals even if they are not capable to do the tasks themselves.

As the randomization region grows, the agent is less biased towards navigating to the upper-right corner, and hence our method’s impact on correcting that bias diminishes. On the other hand, for small randomization regions, the goal is indistinguishable from its position, and hence the behavior of the agent cannot be corrected. Without a change in the training dataset or in the model architecture, the agent will confound the goal. It is interesting to note that regardless of the training method, the agent always learns a location rather a feature. This position-bias could happen for many reasons. First, it could be that learning based on a position may be simpler and more efficient than learning based on some feature. Second, it could be that the goal is not salient to the agent, and the positional information is more prominently represented and easier to interpret. Finally, given that the state space inherently provides more distinct cues about position in the form of coordinates, the learning could naturally prioritize position.