PoEM: Predicting RL Outcomes from Existing Policies
PoEM approximates a reinforcement-learning policy for a new reward from existing post-trained models, without extra RL.
PoEM predicts the policy that reinforcement learning would produce for a new reward using models already post-trained on other rewards, avoiding another full RL run when the reward changes. If the new reward is a linear combination of existing rewards, the new log-policy is a linear combination of the existing log-policies; the authors also observe that log-policies often span an approximately low-rank subspace even when rewards are not linearly related. Combination weights can be estimated from reward or basis-policy outputs on samples. The method is validated on synthetic and real rewards across text and image modalities.
- PoEM predicts a new reward's RL policy without running additional RL training.
- A linear reward mix corresponds to a linear mix of existing log-policies.
- Log-policies often lie in an approximately low-rank subspace across rewards.
- Weights are estimated from reward or basis-policy outputs on samples only.
- Validation covers synthetic and real rewards in text and image modalities.
Full article233 words · extracted from arxiv.org · click to collapse
Foundation models are post-trained with reinforcement learning (RL) to maximize specific rewards, such as human alignment, 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 RL outcomes without actually running RL on it? We answer this in the affirmative by introducing PoEM, a framework to predict the outputs of RL on a new reward function using a set of models already post-trained on other rewards. First, we show that if the new reward function can be written as a linear combination of existing ones, then the new policy in log-space can be written as a linear combination of the existing log-policies. Surprisingly, even in cases where the rewards are not linearly connected, we observe that often log-policies from RL training span an approximately low-rank subspace across rewards. To our benefit, the weighting coefficients 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 new reward function, and approximates the target RL policy without actually running any additional RL training. We experimentally validate our approach across synthetic and real rewards, spanning both text and image modalities.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.30226