OBLR-PO: A Theoretical Framework for Stable Reinforcement Learning

TL;DR

<3-5 sentence high-level summary> The paper addresses instability in RL-based post-training for large language models by developing a unified theoretical framework that characterizes unbiasedness and variance of policy-gradient estimators, derives a loss upper bound, and prescribes an SNR-guided adaptive learning-rate coupled with a gradient-weighted baseline. This leads to the Optimal Baseline and Learning-Rate Policy Optimization (OBLR-PO) algorithm, which jointly adapts learning rates and baselines to improve stability and performance. The authors prove convergence guarantees and demonstrate practical gains on Qwen3-4B-Base and Qwen3-8B-Base across multiple mathematical reasoning benchmarks, validating the theory's applicability. The work bridges principled theory with empirical post-training improvements for large-scale language models.

Abstract

Existing reinforcement learning (RL)-based post-training methods for large language models have advanced rapidly, yet their design has largely been guided by heuristics rather than systematic theoretical principles. This gap limits our understanding of the properties of the gradient estimators and the associated optimization algorithms, thereby constraining opportunities to improve training stability and overall performance. In this work, we provide a unified theoretical framework that characterizes the statistical properties of commonly used policy-gradient estimators under mild assumptions. Our analysis establishes unbiasedness, derives exact variance expressions, and yields an optimization-loss upper bound that enables principled reasoning about learning dynamics. Building on these results, we prove convergence guarantees and derive an adaptive learning-rate schedule governed by the signal-to-noise ratio (SNR) of gradients. We further show that the variance-optimal baseline is a gradient-weighted estimator, offering a new principle for variance reduction and naturally enhancing stability beyond existing methods. These insights motivate Optimal Baseline and Learning-Rate Policy Optimization (OBLR-PO), an algorithm that jointly adapts learning rates and baselines in a theoretically grounded manner. Experiments on Qwen3-4B-Base and Qwen3-8B-Base demonstrate consistent gains over existing policy optimization methods, validating that our theoretical contributions translate into practical improvements in large-scale post-training.

Figures (7)

• Figure 1: The left figure shows the Advantage during training for the Qwen3-4B-Base model, comparing GRPO (blue) and OBLR-PO (orange). The right figure shows the Advantage during training for the Qwen3-8B-Base model, again comparing GRPO (blue) and OBLR-PO (orange).
• Figure 2: The left figure shows the Gradient Norm during training for the Qwen3-4B-Base model, comparing GRPO (blue) and OBLR-PO (orange). The right figure shows the Gradient Norm during training for the Qwen3-8B-Base model, again comparing GRPO (blue) and OBLR-PO (orange).
• Figure 3: The left figure shows the Loss across training steps for Qwen3-4B-Base, comparing GRPO (blue) and OBLR-PO (orange). The right figure shows the Loss across training steps for Qwen3-8B-Base, again comparing GRPO (blue) and OBLR-PO (orange).
• Figure 4: Advantages observed for Qwen3-4B-Base (left) and Qwen3-8B-Base (right) during training.
• Figure 5: Gradient norm curves of Qwen3-4B-Base (left) and Qwen3-8B-Base (right) during training.

Theorems & Definitions (27)

• Theorem 1: Unbiasedness
• Theorem 2: Covariance Expression
• Theorem 3: Upper Bound
• Theorem 4: Optimal Learning Rate Schedule
• Theorem 5
• Theorem 6: Convergence Analysis
• Theorem 7
• Theorem A.1: Unbiasedness
• proof
• Theorem A.2: Variance Expression