Mitigating Reward Hacking via Task Representations

Language models fine-tuned with reinforcement learning frequently learn to exploit their reward source rather than solve the underlying task. Existing fixes are typically reactive: detect a specific exploit, then patch the reward. We instead ask whether reward hacking can be prevented by constraining how the model represents the task. We introduce prompt KL regularization, a single auxiliary loss that adds a KL penalty between the actor's and reference model's per-token distributions on the prompt only, leaving the response distribution free. Across three reward hacking tasks of different sizes, prompt KL regularization maintains low hack rates (under 3%) when standard training yields up to 100% hack rates. Additionally, it matches or improves ground-truth accuracy, with no changes to the reward or environment. We aim to understand why prompt KL regularization is effective. Two independent mechanistic studies suggest that changes in prompt representations are key to models learning reward hacking behavior. First, we show that prompt-space activation steering Pareto-dominates response-space steering on hack rate and accuracy. Then, we demonstrate that swapping prompt activations between models can transfer and reduce reward hacking. Together, these results suggest that the way models internally represent the task is central to reward hacking, and that this representation can be directly constrained to prevent it.