Variance optimal state sampling improves reinforcement learning efficiency
Variance-Optimal State Resampling for Reinforcement Learning with Verifiable Rewards
Machine LearningArtificial Intelligence
Summary
Reinforcement learning often wastes time exploring states that don’t help improve the policy. The authors studied how to pick states to minimize randomness in learning, finding a way to focus on states that matter most for better learning signals. They developed a method called VSR that changes how states are sampled using a clever replay buffer and a low-cost way to estimate variance. This method performs better than similar approaches on challenging tasks while using less computational effort.
What this means in practice
- •For reinforcement learning engineers: Improve training efficiency on complex tasks by reallocating rollout steps to states that yield higher learning signal variance.
- •For game ai developers: Enhance policy learning in games with long horizons like Sokoban or BabyAI by adopting variance-optimal state resampling techniques.
Authors
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, $μ_{\mathrm{var}}^π(s) \propto d^π(s)\,σ_π(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.