PoEM: Predicting RL Outcomes from Existing Policies

1 Introduction

• RL post-training is expensive, sometimes unstable, and must be repeated when rewards change; adapter averaging degrades with more models, while best-of-N is effective only in narrow regimes.The introduction frames the goal as predicting outcomes on new rewards without running RL on them.Source: §1 · p. 1
• When a target reward is a linear combination of basis rewards, its optimum policy has a weighted log-mixture form whose coefficients can be estimated from sampled reward outcomes.This provides a way to simulate RL outcomes on such combinations without running additional RL.Source: §1 · p. 1
• Even when rewards are not linearly connected, single-reward policies often span a low-dimensional log-ratio space that can contain a new optimal policy.This replaces linearity of rewards with a weaker geometric assumption about log-policies.Source: §1 · p. 2
• PoEM estimates composition weights from target-reward samples using basis reward outputs or policy log-ratios, then composes existing policies at inference time without parameter updates or new RL.Figure 1 depicts shared calibration samples, within-prompt centering, and least-squares weight fitting.Source: §1 · p. 2, Figure 1 · p. 2
• 9 of 10 held-out reward-model rewards yield PoEM policies closer to directly trained RL than the leading single expert; text and image experiments also support the approach.On composite text rewards, PoEM generations recover most of the reward gain of a directly trained RL policy; on image generation, PoEM approximates a held-out RL adapter.Source: §1 · p. 2

2 Predicting RL Outcomes by Policy Composition

• PoEM composes existing single-reward policies to approximate the policy that RL training would produce for a new reward.The framework covers linear-reward composition, weight recovery, and checking basis coverage.Source: §2 · p. 3

2.1 Problem setup

• With a shared reference and KL coefficient, the goal is to approximate the policy optimized for a new target reward.The basis policies were obtained by post-training the same reference model on different rewards.Source: §2.1 · p. 3

2.2 Exact composition for linear rewards

• When the target reward is a linear combination of basis rewards, the ideal target policy is a product of experts built from the reference and basis policies.This identity assumes exact optimization of each basis reward; in practice, PoEM composes the implicit rewards represented by trained policies.Source: §2.2 · p. 3

2.3.1 PoEM for Autoregressive language models

• For language models, PoEM applies log-ratio composition locally at each response prefix, using reference and basis-policy next-token logits without parameter updates.Local normalization is not generally identical to the globally normalized sequence distribution.Source: §2.3.1 · pp. 3–4

2.3.2 PoEM for Diffusion models

• For diffusion models, linear reward combinations yield a target denoiser formed from existing denoisers plus a correction term.Dropping the correction is an approximation; it vanishes as noise approaches zero and is zero for basis policies that are Gaussians with shared covariance.Source: §2.3.2 · p. 4

2.4 Recovering composition weights

• PoEM estimates composition weights by regression on centered calibration rewards or, when rewards are unavailable, on basis-policy log-ratios.Reward-space weights recover true weights when the reward combination assumption holds and the calibration matrix has sufficient rank; policy-space weights can differ when optimization or effective strengths differ.Source: §2.4 · p. 5

2.5 Policy space coverage beyond linear rewards

• PoEM can extend beyond linear reward combinations when the target policy lies in the span of basis log-ratios, assessed with held-out explained variance.Log-ratios use fewer effective directions than rewards in the cited 20-expert comparison; high coverage supports using PoEM, while low coverage suggests abstaining or training a new basis policy.Source: §2.5 · pp. 5–6, Figure 2 · p. 6

3 Experimental Setup

• The evaluation tests linear-combination rewards, policy geometry, and held-out rewards against policies trained directly with RL.Each directly trained target policy serves as an oracle for PoEM.Source: §3 · p. 6

3.1 Basis training

• P-GRPO and P-DPO use 20 programmatic text rewards and rank-64 LoRA experts on Qwen3-0.6B, differing in whether experts are trained with GRPO or offline DPO.The programmatic rewards are deterministic response-text functions; both bases share rewards and initialization.Source: §3.1 · p. 7
• The other bases include PPO-trained language experts on reward models and Stable Diffusion experts trained on image statistics or image reward models.RM-Div uses ten diverse public reward models; SD-DDPO has 13 LoRA experts on Stable Diffusion v1.4.Source: §3.1 · p. 7

3.2 Datasets and target policies

• Language experts and targets use UltraChat prompts, with harmful requests included in about half of RM-HH and RM-Div training prompts; evaluation prompts are held out.Programmatic-basis calibration uses held-out UltraChat prompts, whereas RM-basis calibration uses training prompts; target policies follow their basis training recipe.Source: §3.2 · p. 7

3.3 Evaluation and metrics

• Policies are greedily decoded and scored against reference and target policies; composition weights and decoding strengths are set without using the target policy, with best-of-N and top-expert baselines.Programmatic-basis weights are fit on reference responses; held-out and reward-model-basis weights use non-negative least squares on expert samples, excluding the held-out expert.Source: §3.3 · pp. 7–8
• The optional strength γ rescales the composed contribution at test time, compensating for shrinkage when weakly correlated basis directions are averaged.The text notes that GRPO and DPO discard reward scale, unlike PPO with a KL penalty in the reward; it also describes a drift-matched alternative strength.Source: §3.3 · pp. 7–8
• Evaluation measures reward-gain recovery and relative KL to the target policy, while geometry uses explained-variance shares for rewards, log-ratios, and expert weight updates.A relative KL below 1 means the prediction is closer to the target than the reference model is; the geometry calculation uses 2,400 reference responses to 600 prompts.Source: §3.3 · p. 8
• Held-out-reward coverage is estimated from expert-sampled calibration responses, with a 0.3 threshold except for RM-Div, whose values all fall below it.For RM-Div, the authors instead split rewards at median coverage and use within-expert coverage.Source: §3.3 · p. 8

4.1 Combined rewards

• 0.83–1.00 of πtar’s reward gain is recovered on programmatic bases after γgeom; RM bases recover 0.78–1.08 at γ = 1.On programmatic bases, PoEM’s recovery error is 0.19–0.28, close to the 0.12–0.17 variation between RL runs.Source: Figure 3 · p. 9, §4.1 · p. 9
• On P-GRPO and RM bases, PoEM is closer to πtar than the reference on all combined rewards; P-DPO matches reward or distance, but not both.On RM combined rewards, PoEM is closer to πtar than the top expert and, unlike best-of-N, does not query the target reward at decoding.Source: §4.1 · p. 9

4.2 Beyond the linear case: policy space geometry of post-trained models

• A target reward need not be a linear combination of basis rewards if its optimal log-policy lies near the span of basis log-ratios.This weaker geometric condition also makes the regression’s effective dimension the policy-space rank, so a small calibration set can suffice.Source: §4.2 · p. 10
• Centered eΦ has effective rank 6.3 of 20, while adapter updates are full rank at 19.4, despite near-orthogonal updates with median pairwise cosine 0.02.The adapters’ induced log-policies have fewer effective directions than their parameter updates, even for rewards without an obvious linear relationship.Source: §4.2 · p. 10

4.3 Held-out rewards

• 0.55–0.69 of RL’s reward gain is recovered for covered held-out rewards, versus 0.22–0.43 for uncovered rewards; coverage ranks reward error at ρ = −0.71 to −0.81.Coverage ranks rewards within a basis, not across bases: RM-RB has near-zero coverage despite 0.56–0.83 recovery, as its reward models largely agree.Source: §4.3 · p. 10, Table 2 · p. 10
• Increasing γ helps covered rewards but cannot supply a missing direction; even γ tuned on πtar leaves uncovered rewards below covered ones on P-GRPO and P-DPO.On P-GRPO, increasing γ from 0.5 to 2 raises median covered recovery from 0.18 to 0.64, while the median uncovered prediction reaches only 0.26 and drifts away.Source: §4.3 · p. 11

4.4 Image diffusion models: qualitative results

• Composing experts reproduces held-out image-reward effects, and the image-model composition fits weights separately at each denoising step.The reported effects include more saturated colors, finer texture, and a centered subject on a blurred background; the held-out expert’s effect lies largely within what the other experts express.Source: §4.4 · p. 11

Policy composition

• Parameter averaging works when fine-tuned checkpoints share a basin and linear mode connectivity holds, but its benefits decline as more models are combined.The passage attributes this decline to updates in different directions diluting each other.Source: §5 · p. 11
• DExperts, contrastive decoding, proxy-tuning [29], and MOD mix logits at decoding; MOD uses user-chosen weights.PoEM’s composition rule follows the closed-form KL-regularized RL solution, and it fits unknown weights from experts’ log-ratios on a small calibration set.Source: §5 · p. 12

Multi-objective alignment

• Multi-objective alignment methods choose reward weights during training, whereas PoEM composes existing single-reward experts at decoding and fits weights afterward.This means weights do not have to be chosen before training.Source: §5 · p. 12

Implicit rewards and policy space structure

• Prior work uses policy log-ratios as response scores for self-training, preference-pair selection, and process reward models; PoEM regresses the target reward on multiple experts’ log-ratios.Work on datamodels and linear mode connectivity also finds fewer dimensions in policy space than parameter space.Source: §5 · p. 12

What RL post-training changes

• Although expert weight updates are nearly orthogonal, their policies move along a few shared directions; the authors suggest this explains why composition can predict RL when the target direction is represented.They note that this has only been checked on small models.Source: §6 · p. 12

Previewing a reward before training on it

• PoEM can preview a candidate reward’s effect without training, but the preview should be trusted only when coverage is high.Source: §6 · p. 13

Approximating expensive rewards

• PoEM uses reward evaluations on a small calibration set rather than every rollout, and low coverage means it fails when an expensive reward depends on an unsupported direction.The discussion describes RM-Div predicting held-out 7–8B reward models from experts trained on the others.Source: §6 · p. 13

Building the basis

• Coverage can identify rewards the basis cannot reach, while inference with k experts is costly; proposed distillation and RL-start alternatives remain untested.Source: §6 · p. 13

• For rewards outside the basis combinations, coverage ranks which targets are reachable; low-coverage rewards remain far from the RL outcome at every tested γ.Source: §7 · p. 13
• On reward-model combinations, PoEM is closer to the RL outcome than both the reference model and the top expert, while coverage ranks held-out reward models.Source: §7 · p. 13
• For linear-combination rewards, composing experts at a scale set without new training achieves reward performance similar to a second RL run.Source: §7 · p. 13

Limitations

• Coverage cannot recover target directions outside the chosen policy-space span, and exact composition assumes KL-regularized experts sharing a reference model and β.The authors also report not testing long-context generation or transfer of fitted weights between reference models.Source: Limitations · p. 13

A A guarantee under coverage

• A uniform coverage condition bounds how far the sequence-level product can be from the target optimum; the proof uses reward-shift invariance and uniform ε-closeness.Source: A A guarantee under coverage · p. 18

B Experimental details

• Appendix B provides the models, data, and protocol details omitted from Section 3.Source: B Experimental details · p. 18

B.1 Programmatic bases

• The programmatic setup defines 20 rewards and 32 combined targets across k=2, 4, 8, 10, and 16, using hand-set or Dirichlet-drawn weights.Source: §B.1 · pp. 18–19

B.2 Reward model bases

• RM-RB, RM-HH, and RM-Div use reward-model bases; RM-Div’s ten members were fixed before held-out decoding following health checks.Four members retained mild surface tics, and the OASST-Pythia expert complied more with harmful requests than the others.Source: §B.2 · pp. 19–20

B.3 Diffusion basis

• For diffusion experts, each denoising step fits the held-out expert’s noise prediction by least squares with an intercept using the other twelve experts at that latent.Source: §B.3 · p. 20

B.4 Evaluation

• Evaluation uses calibration prompts from the held-out UltraChat pool, drops 48 prompts overlapping the evaluation set for weight and coverage fits, and uses all 452 for the γgeom Gram matrix.Source: §B.4 · pp. 20–21
• For programmatic combined rewards, ridge regression selects weights with cross-validation; coverage averages validation R2 across 20 splits, each holding out 20% of prompts.Source: §B.4 · p. 21
• On reward-model bases, drift-matched γ⋆ is selected against a non-decreasing drift curve built from PoEM decodes at γ values 0.5 through 3, with γ=0 assigned zero drift.If the target drift exceeds the curve, γ⋆ is set to 3; obtaining it requires decoding each expert and PoEM at every grid value.Source: §B.4 · p. 21

B.5 Geometry

• The geometry analysis uses 600 prompts, computes ΔΘ only for bases sharing one initialization, and reports bootstrap intervals from 200 prompt-level draws.Reward and log-ratio predictions use ordinary least squares on five prompt-disjoint 80/20 splits.Source: §B.5 · p. 21

C.1 Additional results on combined rewards

• On combined reward-model rewards, PoEM ends closest to the directly trained target policy across RewardBench and ArmoRM combinations.The figures compare reward-gain recovery and KL divergence for PoEM and other methods.Source: Figure 7 · p. 22
• On nine P-GRPO combined rewards, PoEM’s median reward error at γgeom matches the second RL run at 0.12, though that run is closer on six rewards.Retraining with GDPO misses the first run by 0.09 on six P-GRPO rewards; on P-DPO, the second run is only slightly closer than the reference model by KL (0.54 versus 0.69).Source: §C.1 · pp. 21–22
• For k ≥ 4, γgeom closes about half the recovery gap on P-DPO and most of it on P-GRPO; it keeps recovery around 0.8–1.2 as more rewards are composed.At γ = 1, PoEM’s recovery drops once more than two rewards are composed.Source: Figure 8 · p. 23, Figure 9 · p. 23

C.2 Gap between the decoder and the sequence-level product

• The sequence-level product differs from the token-normalized decoder by 0.6–2.7 nats per response, but reweighting changes the combined reward by at most 0.03 z-units.The sequence-level product is further from πtar than the decoder in all six settings, significantly in five; the authors conclude local normalization is not what limits PoEM on this basis.Source: Table 9 · p. 25, §C.2 · p. 25

C.3 Held-out reward recovery

• On P-GRPO, every covered held-out reward ends closer to πtar than πref at γ = 1, γgeom, and γ⋆; uncovered rewards retain large error even with γ tuned on πtar.Figure 11 relates coverage to recovery and distance to πtar, while Figure 12 contrasts covered and uncovered rewards.Source: Figure 11 · p. 25, Figure 12 · p. 26

C.4 Policy space geometry: additional figures

• The ten public reward models and both programmatic bases show a gap between policy-update and log-ratio geometry, with ratios 0.50 and 0.75 for the programmatic bases.Figure 13 reports cumulative variance for weight updates, rewards, and log-ratios.Source: Figure 13 · p. 26

C.5 Experts trained with reward models: additional results

• On RM-Div, the ten log-ratios have effective rank 2.84 on reference-model responses, versus 6.96 for the ten rewards; on RM-RB, the corresponding ranks are 1.64 versus 2.30.On experts’ own responses, RM-Div ranks are 6.09 versus 8.00, while RM-RB log-ratio rank (2.77) exceeds reward rank (2.43).Source: §C.5 · p. 26
• Within-expert coverage ranks held-out RM-Div recovery better than per-prompt coverage (Spearman ρ = 0.81 versus 0.45), while RM-RB recovery reaches 0.56–0.83 despite both scores below 0.12.Within-expert coverage also ranks relative KL better on RM-Div (ρ = −0.49 versus −0.16).Source: §C.5 · p. 26
• On RM-Div, γ⋆ raises median recovery from 0.60 to 0.66 for the more-covered half and from 0.43 to 0.51 for the less-covered half; on RM-RB it changes recovery from 0.79 to 0.81.For the less-covered RM-Div half, tuning γ on πtar raises recovery to 0.67 but yields relative KL of 2.04.Source: §C.5 · p. 27
• For OASST-Pythia-1.4B and InternLM2-7B, the other log-ratios explain 82% and 92% of each expert’s log-ratio despite explaining only 21% and 63% of its reward variance.At γ = 1, PoEM recovers 0.43 and 0.57 of their gains, respectively.Source: §C.5 · p. 27, Figure 14 · p. 27

C.6 Diffusion models: additional results

• The image recovery measure ignores image quality and likeness to the held-out expert’s images, and span R2 above 0.99 for every reward does not track recovery (Spearman ρ = 0.12).Recovery counts only the gain in the held-out reward; span R2 uses the full noise prediction rather than each expert’s change to the reference prediction.Source: §C.6 · p. 27, Table 10 · p. 28
• Composing the other twelve experts recovers more than half the held-out expert’s reward gain on 9 of 13 image rewards, averaging 0.67 across rewards.Recovery ranges from 0.25 on aesthetic score to 1.07 on CLIP score; the weights are fit at each denoising step to the held-out expert’s noise prediction, so this is not prediction without that expert.Source: §C.6 · p. 27, Table 10 · p. 28
• Across the displayed diffusion rewards, most vary independently; only the image-detail group moves together, and three listed rewards have no expert in SD-DDPO.The group includes sharpness, edge density, entropy, and two JPEG rewards; brightness, hue diversity, and rule of thirds have no expert in SD-DDPO.Source: Figure 15 · p. 28