Variance-Optimal State Resampling for Reinforcement Learning with Verifiable Rewards
Yibo Wang ⋅ Yuanxin Liu ⋅ Zhaoran Wang
Abstract
Reinforcement learning with verifiable rewards spends much of its rollout budget on states that yield no learning signal. We recast this as a question about the sampling distribution itself: under a fixed rollout budget, which distribution over states minimizes the variance of the policy-gradient estimator? The answer admits a closed-form Neyman allocation, $\mu_{\mathrm{var}}^\pi(s) \propto d^\pi(s)\,\sigma_\pi(s)$, which connects naturally to GRPO's group-relative advantage and induces an emergent reverse curriculum as the policy improves, with no externally imposed schedule. Building on this principle, we propose VSR (Variance-driven State Resampling), a plug-in modification to GRPO using a prefix replay buffer and a zero-cost variance proxy. On mathematical reasoning benchmarks and long-horizon agentic tasks (Sokoban, BabyAI), VSR outperforms GRPO and GiGPO on most benchmarks with fewer generated tokens per update.
Chat is not available.
Successful Page Load