Start Classifying: Categorical Critics for LLM Reinforcement Learning

2026-08-03

AI Summary

In reinforcement learning with verifiable rewards (RLVR), sparse binary rewards render conventional PPO critics trained with mean squared error highly sensitive to value estimation errors, leading to distorted advantage estimates. This work proposes HL-Gauss PPO, the first method to integrate a categorical critic into PPO training for large language models: it performs classification over a discrete value support using cross-entropy loss and decodes the expected value for standard generalized advantage estimation (GAE) and policy updates. Without altering the policy optimization mechanism, this approach substantially improves the calibration and stability of critic signals. Evaluated on Qwen2.5 and Qwen3 models across mathematical reasoning, tool-augmented math, and Search-R1 tasks, HL-Gauss PPO consistently outperforms PPO and DAPO baselines, yielding lower-variance, more symmetric advantage distributions and significantly reduced Brier scores and calibration errors.

Abstract

Proximal Policy Optimization (PPO) for large language models typically trains its critic by mean-squared-error (MSE) regression on scalar value targets. Although scalar MSE is statistically valid for estimating the conditional expected return, sparse binary rewards in reinforcement learning with verifiable rewards (RLVR) make critic optimization and calibration especially consequential: small value errors directly distort the scalar advantages used by PPO. We study whether a classification-based training objective can improve this critic signal. HL-Gauss PPO replaces the scalar MSE head with a categorical predictor over a discretized value support, trained by cross-entropy against smoothed HL-Gauss targets. Its output is decoded to a scalar expectation for standard GAE and PPO; the actor update is therefore unchanged and is not distributional. Across mathematical reasoning, tool-augmented math, and Search-R1, and on both Qwen2.5 and Qwen3 backbones, HL-Gauss PPO consistently improves over strong PPO and DAPO baselines. Controls with one-hot, two-hot, and Bernoulli two-bin critics show that neither a larger output head nor binary classification alone explains the gains. On a common collection of reasoning prefixes, HL-Gauss improves Brier score and calibration error and yields more symmetric, lower-variance advantages. These results position categorical value learning as an effective optimization surrogate for PPO critics in RLVR.

Problem

• reinforcement learning with verifiable rewards
• critic optimization
• sparse binary rewards
• value estimation error
• advantage distortion

Innovation

• categorical critics
• HL-Gauss PPO
• reinforcement learning with verifiable rewards
• value function calibration
• distributional reinforcement learning