Scale equalized rollout allocation improves reinforcement learning gradient estimates
SERA: Scale-Equalized Rollout Allocation for Maximum Likelihood Reinforcement Learning
Machine LearningArtificial Intelligence
Summary
Reinforcement learning methods sometimes suffer from uneven attention to different cases when they have limited attempts to learn from. The authors show that this unevenness can weaken learning for less successful cases, skewing overall progress. They propose a way to balance the number of attempts each case gets so that the learning signals are fairer and more accurate. This method better matches ideal learning updates and improves performance on tasks like maze solving and math reasoning.
What this means in practice
- •For ai research teams: Improve training accuracy and stability of reinforcement learning models for reasoning problems by balancing computational effort across different data prompts.
- •For robotics software developers: Enhance robot navigation algorithms by better distributing simulation runs across varied scenarios for improved policy learning efficiency.
Authors
Zihao Chen, Fanxiang Xiong, Hongran Ren, Xuefeng Bai, Zhongxiang Dai, Kehai Chen, Zhiguo Zhang, Zhiyong Wang, Yu Cheng
Abstract
Maximum Likelihood Reinforcement Learning (MaxRL) targets prompt-wise log-success and has shown strong performance on reasoning tasks. Under finite rollout budgets, however, the estimator used by MaxRL attenuates each prompt's likelihood gradient by a factor that depends on its success probability and rollout count. Under uniform rollout allocation, the common rollout count fails to compensate for success-dependent attenuation, leaving low-success prompts more strongly attenuated and distorting their relative contributions to the expected aggregate gradient. We introduce SERA (Scale-Equalized Rollout Allocation), which redistributes a fixed rollout budget to approximately equalize these finite-rollout scaling factors. Building on our theoretical analysis of how finite rollouts distort prompt-wise likelihood gradients, we formulate the allocation as a fixed-budget max--min problem, derive a waterline solution to its continuous relaxation, and introduce a multiplicity correction to remove the additional prompt weighting induced by heterogeneous rollout counts. Experiments show stronger alignment with exact likelihood gradients in a controlled ImageNet setting and improved multi-sample solution coverage over MaxRL on maze navigation and mathematical reasoning under matched training rollout budgets.