Source-linked AI summary
Truncated Proximal Policy Optimization
Tiantian Fan, Lingjun Liu, Yu Yue, Jiaze Chen, Chengyi Wang, Qiying Yu, Chi Zhang, Zhiqi Lin, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Bole Ma, Mofan Zhang, Gaohong Liu, Ru Zhang, Haotian Zhou, Cong Xie, Ruidong Zhu, Zhi Zhang, Xin Liu, Mingxuan Wang, Lin Yan, Yonghui Wu
TL;DR
PPO’s on-policy nature limits efficiency when training on long chain-of-thought trajectories. T-PPO updates policy from incomplete responses and reports up to 2.5× training speedup with competitive final performance.
Problem
PPO’s on-policy nature limits training efficiency on long chain-of-thought trajectories, which can entail substantial computational overhead and extended training durations.
Source: §1 · p. 2Method
Results
Up to 2.5× training speedup with competitive final performance is reported for T-PPO.
Source: §5 · p. 9Takeaways & Limitations
The authors present T-PPO’s empirical findings as practical guidance for research on efficient large-scale reinforcement learning.
Source: §5 · p. 9EGAE assumes that the value of the unseen next state in a truncated trajectory equals that of the last generated state.
Source: §3.1 · p. 5
Abstract
from arXiv · showhide
Recently, test-time scaling Large Language Models (LLMs) have demonstrated exceptional reasoning capabilities across scientific and professional tasks by generating long chains-of-thought (CoT). As a crucial component for developing these reasoning models, reinforcement learning (RL), exemplified by Proximal Policy Optimization (PPO) and its variants, allows models to learn through trial and error. However, PPO can be time-consuming due to its inherent on-policy nature, which is further exacerbated by increasing response lengths. In this work, we propose Truncated Proximal Policy Optimization (T-PPO), a novel extension to PPO that improves training efficiency by streamlining policy update and length-restricted response generation. T-PPO mitigates the issue of low hardware utilization, an inherent drawback of fully synchronized long-generation procedures, where resources often sit idle during the waiting periods for complete rollouts. Our contributions are two-folds. First, we propose Extended Generalized Advantage Estimation (EGAE) for advantage estimation derived from incomplete responses while maintaining the integrity of policy learning. Second, we devise a computationally optimized mechanism that allows for the independent optimization of the policy and value models. By selectively filtering prompt and truncated tokens, this mechanism reduces redundant computations and accelerates the training process without sacrificing convergence performance. We demonstrate the effectiveness and efficacy of T-PPO on AIME 2024 with a 32B base model. The experimental results show that T-PPO improves the training efficiency of reasoning LLMs by up to 2.5x and outperforms its existing competitors.
1 Introduction
T-PPO improves the efficiency of on-policy reasoning-model training by updating policies from partial responses while retaining unbiased value estimation from completed rollouts. On AIME 2024, it reports 62 pass@1 and 2.5× higher training efficiency than state-of-the-art synchronization algorithms.
1 Introduction
- PPO [4] is on-policy, making long chain-of-thought trajectories costly to process and extending training time.Off-policy PPO variants seek greater sample efficiency through trajectory reuse.Source: §1 · p. 2
- GePPO extends policy-improvement guarantees off-policy, Off-Policy PPO uses off-policy data to avoid large updates, PPO-EWMA uses decoupled objectives and EWMA, and KIMI K1.5 reuses previous trajectories; off-policy methods typically suffer high policy-gradient variance, unstable training, and degraded performance.Source: §1 · p. 2
- EGAE enables policy updates from partially generated responses, while Monte Carlo value updates wait for complete rollouts and use observed returns to avoid value-estimation bias.Selective token screening supports simultaneous but independent policy and value model optimization.Source: §1 · p. 2
- 62 pass@1 on AIME’24 and 2.5× higher training efficiency than state-of-the-art synchronization algorithms are reported for T-PPO.The paper also reports robust convergence without additional constraints or regularization beyond standard PPO.Source: §1 · p. 2
2 Preliminary
The preliminaries define the sparse-reward generation setting, review PPO and GAE, and explain T-PPO’s use of Monte Carlo value estimates for reasoning tasks.
2.1 Reinforcement learning framework
- LLM generation is modeled as an MDP whose state is the prompt plus generated prefix; when correctness reward occurs only at the terminal step, the task reduces to a bandit formulation.The paper notes that its approach can also apply to simpler settings with process rewards.Source: §2.1 · p. 3
2.2 Proximal Policy Optimization
- PPO is a default LLM actor-critic baseline whose clipped objective stabilizes policy updates, but on-policy sample collection between updates makes it inefficient.The clipping objective limits policy changes to preserve training stability, while high-variance updates require many samples collected under the current policy.Source: §2.2 · p. 3
2.3 Generalized Advantage Estimation
- GAE [9] trades bias against variance by exponentially weighting multi-step advantage estimates, with λ controlling their relative weights.The TD residuals also depend on the discount factor γ, which determines the value assigned to future versus immediate rewards.Source: §2.3 · pp. 3–4, §2.2 · p. 3
2.4 Value Function Estimation
- The critic estimates returns and supplies learning signals; for reasoning tasks, T-PPO uses Monte Carlo (TD(1)) value estimates to avoid GAE’s temporal-difference bias.This choice maintains strictly unbiased state-value predictions despite forgoing GAE’s variance-reduction benefits.Source: §2.4 · p. 4
3 T-PPO: Truncated Proximal Policy Optimization
T-PPO improves PPO efficiency by optimizing incomplete trajectories. It combines EGAE, truncated rollouts with selective token use, and independent simultaneous policy and value updates.
3.1 EGAE: Extended Generalized Advantage Estimation
- EGAE estimates advantages from unfinished trajectories by replacing the full horizon T with truncation length l, assuming V(sl) = V(sl−1).Its discounted Bellman-residual formulation retains GAE’s bias-variance trade-off, with future reward influence controlled by λ.Source: §3.1 · pp. 4–5, §3 · p. 4
3.2 Token Filtering Strategy
- Truncated rollouts replace completed sequences with new prompts while retaining unfinished samples; policy training uses current-step response tokens, while value training uses finished-sequence tokens.Some latest tokens from unfinished sequences may be excluded from policy training because their advantage estimates can have high variance.Source: §3 · p. 4, §3.2 · p. 5, Figure 2 · p. 5
3.3 Implementation and Training Efficiency Analysis
- L/l = k implies approximately k-fold generation and training-time savings, and about k-fold end-to-end efficiency for the same number of training steps.The algorithm collects rollouts for window length l, estimates advantages with EGAE, then updates the policy and value function.Source: §3.3 · p. 6
4 Experiments
On AIME 24, T-PPO reaches 61.88 pass@1 while matching PPO with 20k response length at 60% wall-clock time reduction; its experiments also examine convergence efficiency and response-length dynamics.
4.1.1 Models and Configurations
- T-PPO uses Qwen-2.5-Base-32B, 24k maximum responses, and 8k windows, with KL divergence removed to encourage exploration.The policy and critic learning rates are 1e-6 and 2e-6, respectively; evaluation reports avg@32, with temperature 1.0 and topp 0.7.Source: §4.1.1 · pp. 6–7
4.2.1 Main Results
- 61.88 pass@1 on AIME 24 surpasses DeepSeek-R1-Zero-Qwen-32B and the available SOTA async PPO, while matching PPO with 20k response length at 60% wall-clock time reduction.The paper reports these as the main results for T-PPO.Source: Table 1 · p. 7, §4.2.1 · p. 7
4.1.2 Dataset Description
- Experiments compare vanilla PPO, asynchronous PPO-EWMA, and T-PPO on AIME, using DAPO-Math-17K training data and a correctness-based Math-Verify reward.The test set comprises problems from the last year; the training set supplements past AIME questions with artificially constructed difficult math problems. The reward is 1 for a correct final answer and 0 otherwise.Source: §4.1.2 · p. 7
4.2.2 Training Efficiency
- 6720 vs. 11200 convergence steps each: T-PPO and PPO-EWMA have similar per-step wall-clock time, but T-PPO converges in fewer steps.Per 1000 steps, T-PPO's average wall-clock time is comparable to PPO-EWMA's and much lower than vanilla PPO's.Source: §4.2.2 · p. 7, Figure 4 · p. 8
- 249 operations/byte in policy rollout for T-PPO versus PPO's 84 operations/byte places T-PPO closer to the Roofline arithmetic peak.The paper attributes this to improved compute-resource utilization and fewer memory-bound bottlenecks during generation.Source: §4.2.2 · pp. 7–8, Figure 5 · p. 8
4.2.3 Training Dynamics
- Response length initially rises, temporarily declines, recovers, and stabilizes above vanilla PPO, indicating that T-PPO preserves response-length scaling capacity.The paper relates initial length expansion to the emergence of lengthy CoT through RL [14].Source: Figure 6 · p. 8, §4.2.3 · p. 8
5 Conclusion
T-PPO combines successive batching, extended generalized advantage estimation, and efficient policy- and value-model optimization to improve training efficiency while maintaining competitive final performance.
5 Conclusion
- T-PPO extends PPO with successive batching and extended generalized advantage estimation, alongside computationally efficient mechanisms for optimizing policy and value models.Source: §5 · pp. 8–9
- T-PPO achieves up to 2.5× training speedup with competitive final performance.The conclusion also recommends further exploration of truncated or more advanced RL methods.Source: §5 · p. 9
- The experiments and empirical insights provide practical guidance for future research on efficient large-scale reinforcement learning.Source: §5 · p. 9