Reward-Proportional Policy Optimization in Large Combinatorial Spaces
Abstract
Many reinforcement learning problems require learning a distribution over diverse high-quality outcomes rather than converging to a single optimum. Inverse Probability Scaling (IPS) enables reward-proportional sampling by reweighting terminal rewards using outcome probabilities. However, existing IPS approaches estimate these probabilities from empirical within-group frequencies, which become uninformative in large combinatorial spaces where terminal repetitions are rare. We introduce Multiplicity-Aware Inverse Probability Scaling (MIPS), a scalable method estimating terminal probabilities by combining sampled-trajectory likelihoods with a learned multiplicity-allocation model approximating trajectory probabilities conditional on terminal outcomes. The resulting reward scaling integrates directly with Group Relative Policy Optimization (GRPO), preserving the standard policy-gradient objective and avoiding specialized flow-matching losses. Across hypergrid, molecular synthesis, and phylogenetic tree construction tasks, MIPS-GRPO improves sampling performance and diversity over GRPO and IPS-GRPO while achieving performance better than GFlowNet-based methods.