Actor critic method improves training efficiency of large language models
Trust the Critic More
Machine LearningArtificial Intelligence
Summary
Training large language models using reinforcement learning often requires waiting until a whole task is done to understand how well the model performed. This can be slow and inefficient. The authors developed a method called Actor-Critic with Action Chunking that breaks down tasks into smaller chunks and uses a learned critic to evaluate progress without finishing the entire task. Their approach speeds up training and reduces computing resources needed by trusting the critic more under specific conditions.
What this means in practice
- •For machine learning engineers: Train large language models more efficiently by using action chunking and reliable critic updates to reduce the need for full-rollout evaluations.
- •For automated theorem proving teams: Improve proof generation systems by employing actor-critic methods that assign credit to partial proof steps, speeding up model optimization.
Authors
Kaiyue Wen, Luke Bailey, Arvind Mahankali, Tengyu Ma
Abstract
Standard language model RL algorithms credit every token of a long rollout with the same advantage determined by the terminal reward. Actor-critic methods can provide finer-grained credit assignment, but learned critics are generally considered too inaccurate to trust when training LLMs with RL. In recent works, even when a critic is present, it is used only for baseline estimation, so every trajectory must be rolled out to its terminal reward. We introduce Actor-Critic with Action Chunking (AC2) that removes the need to roll every trajectory to completion. AC2 instead assigns credit to action chunks: short continuations of prefixes of past trajectories. A learned critic scores the state reached at the end of each action chunk, allowing the policy to update without observing a terminal reward. We make critic-based credit assignment reliable through three design choices. First, we introduce local readiness which uses critic-based updates on a problem only when the critic is sufficiently accurate on that particular problem. Second, when available, we provide the critic with a reference solution from a previous successful rollout. Third, we assign credit over action chunks of 10k tokens rather than individual tokens, giving the critic a more meaningful portion of the trajectory to evaluate. We train Qwen3-4B on FineProofs-RL using AC2 and evaluate on IMO-ProofBench. AC2 exceeds GRPO's peak validation score of 18.5% using 2.5x fewer decoding FLOPs. This gain comes from two sources, (1) AC2 requires 25% fewer training steps to reach this score, and (2) each step generates fewer tokens because the policy does not need to continue every trajectory to completion. Conceptually, we demonstrate that we can remove the need to roll out every trajectory to completion, opening up a large previously unexplored design space for LLM RL algorithms.