DISA: Offline Importance Sampling for Distribution-Matching LLM-RL

2026-05-17

AI Summary

This work addresses the issue of policy collapse in distribution-matching reinforcement learning, which arises from the coupling between online estimation of the partition function and policy updates, leading to calibration errors. To resolve this, the authors propose decoupling partition function estimation from the reinforcement learning loop by precomputing trajectory-space partition functions offline via importance sampling conditioned on prompts, and freezing this estimate before policy optimization. This approach fully disentangles calibration from learning at the levels of data, gradients, loss, and diagnostics, substantially improving training stability and interpretability. Evaluated on six mathematical and three code generation benchmarks, the method outperforms strong baselines—including FlowRL, GRPO, and GSPO—with a maximum improvement of 13.8 points in Mean@8 while effectively preserving policy diversity.

Abstract

Modern reasoning agents are increasingly evaluated on their ability to generate multiple valid solution paths, plans, or tool-use traces for a given input. Standard reward-maximizing RL tends to collapse onto the most easily reinforced high-reward mode, whereas distribution-matching RL aims to allocate probability mass across the entire reward-shaped solution set. Achieving this objective requires computing a prompt-dependent partition function over the trajectory space. Because existing distribution-matching methods learn this partition function online alongside the policy, calibration errors in the partition function directly distort policy updates and remain impossible to diagnose independently. We introduce DISA, short for Decoupled Importance-Sampled Anchoring, which moves this calibration problem outside the RL loop. DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate before policy optimization begins. This decoupling preserves the distribution-matching objective while strictly separating partition-function estimation from policy learning in data, gradients, loss, and diagnostics. Empirically, on two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline FlowRL, outperforms rewardmaximization baselines GRPO and GSPO on math averages, and exceeds LoRASFT distillation by up to 13.8 Mean@8 points on the same offline trajectories. An LLM-as-judge evaluation further shows that DISA retains substantially more strategy-level diversity than reward-maximization baselines, and sensitivity studies on the proposal strength and inverse temperature follow the bias-variance pattern predicted by the analysis.

Problem

• distribution-matching RL
• partition function
• offline importance sampling
• calibration error
• trajectory diversity

Innovation

• distribution-matching RL
• offline importance sampling
• partition function estimation
• decoupled policy learning
• trajectory diversity