PoEM: Predicting RL Outcomes from Existing Policies

원문 초록

Foundation models are post-trained with reinforcement learning (RL) to maximize specific rewards, such as helpfulness, correctness, or instruction following. This post-training process is computationally intensive, sometimes unstable, and has to be run from scratch every time the reward model changes or when we want to combine multiple rewards. We hence ask: given a new reward function, is it possible to predict the policy RL would produce without actually training on that reward? We answer this in the affirmative by introducing PoEM, a framework for predicting RL outcomes using models previously trained on other rewards. First, we build on the fact that if the new reward function can be written as a linear combination of existing ones, then the optimal policy can be expressed as a combination of the existing policies in log-probability space. Surprisingly, even in cases where the rewards are not linearly connected, we observe that policies learned through RL often span an approximately low-rank subspace across rewards. To our benefit, the weights for this combination can be estimated using only the reward or basis policy outputs on the samples. We turn these observations into an algorithm that takes post-trained models and a target reward function, and approximates the target RL policy without any additional RL training. We experimentally evaluate our approach across programmatic and learned rewards for open-ended generation, with additional experiments on math reasoning, tool calling, and image generation.