PoEM: Predicting RL Outcomes from Existing Policies

TL;DR

PoEM enables prediction of reinforcement learning outcomes for new reward functions without running additional RL training. It leverages existing post-trained models and exploits the low-rank structure of log-policies across rewards to approximate new policies via linear combinations.

Why It Matters

This reduces the need for expensive and unstable RL training cycles when reward functions change or are combined. It accelerates policy adaptation in applications like instruction following and alignment, especially in multimodal settings.

Key Findings

• Log-policies from RL training across different rewards lie in a low-rank subspace, enabling efficient approximation.
• When new rewards are linear combinations of existing ones, the corresponding log-policy can be expressed as a linear combination of existing log-policies.
• Coefficients for the combination can be estimated using only reward scores or policy outputs on sample data, avoiding full RL retraining.