PaperlayerView exact revision on arXiv

Source-linked AI summary

Trust the Critic More

Kaiyue Wen, Luke Bailey, Arvind Mahankali, Tengyu Ma

arXiv:2609.39247v2cs.LGcs.AI

TL;DR

Standard LLM RL assigns a rollout’s terminal-reward advantage to every token, while learned critics are generally used only as baselines. AC2 updates from critic-scored action chunks when the critic is ready for a problem, and exceeds GRPO’s peak mean score with fewer decoding FLOPs on IMO-ProofBench.

  • Problem

    Standard RL gives every token in a rollout the same terminal-reward-based advantage, while learned critics are generally considered too inaccurate for fine-grained credit assignment without terminal rewards.

    Source: §1 · p. 2

  • Method

    AC2 scores continuations of replayed prefixes at action-chunk endpoints and uses critic-based updates only when local readiness indicates sufficient problem-specific accuracy.

    Source: §1 · p. 2

  • Results

    AC2 reached a peak mean score of 20.57%, versus 18.50% for GRPO, using 2.5× fewer decoding FLOPs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    AC2 demonstrates that accurate learned critics can remove the need to observe terminal rewards on every update, opening further design space for LLM RL.

    Source: §6 · p. 11

    Local readiness requires epoching training data to obtain complete rollouts for ground-truth value estimates, which is unrealistic in the non-data-constrained setting.

    Source: §6 · p. 12

Abstract

from arXiv · show

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.

1 Introduction

Standard RL assigns every token in a rollout the same terminal-reward-based advantage, while AC2 uses critic-scored action chunks to update without completing every trajectory. On IMO-ProofBench, AC2 exceeds GRPO’s peak mean score with fewer decoding FLOPs.

1 Introduction

  • AC2 scores action-chunk endpoints with a learned critic, allowing policy updates without observing a terminal reward.Action chunks are short continuations of prefixes from past trajectories.Source: §1 · p. 2
  • AC2 assigns one advantage per action chunk rather than per full trajectory, providing finer-grained supervision while generating fewer tokens per step.The critic’s accuracy is the condition for avoiding terminal rewards on an update.Source: §1 · pp. 2–3
  • Local readiness gates critic-based updates by problem-specific accuracy; reference solutions and 10k-token chunks further support reliable critic estimates.Without local readiness, runs plateau early; shrinking chunks from 10k to 2k tokens also leads to earlier plateaus, while a correct reference lowers critic error.Source: §1 · pp. 2–3
  • 20.57% peak mean score for AC2 versus 18.50% for GRPO, using 2.5× fewer decoding FLOPs.AC2 reached GRPO’s peak in 90 rather than 120 training steps and generated fewer tokens per update by not completing every trajectory.Source: §1 · p. 2
  • Two comparable AC2 variants broaden the design space: sampling one chunk per critic-ready problem and using a highly stale replay buffer for prefixes.The paper presents these variants as examples of algorithms that can leverage existing trajectories more extensively.Source: §1 · p. 3

2 Preliminaries

The paper formulates reinforcement learning with verifiable rewards as maximizing expected scores from a deterministic verifier, with reward available only for complete responses. It contrasts terminal-reward advantage estimates with critic-based estimates that can omit terminal rewards.

2 Preliminaries

  • A deterministic verifier assigns a terminal reward to each complete response, and the learning objective is expected verifier score.The reward is zero before the final token, and the verifier is defined only on complete responses.Source: §2 · p. 3
  • With deterministic transitions and terminal-only rewards, a chunk’s action value is the expected value of the state reached after that chunk.The paper notes that the same identities hold when an action is a fixed-size token chunk.Source: §2 · p. 3
  • GAE’s λ controls the trade-off between learned-value estimates and realized terminal rewards: λ = 0 uses values without terminal reward, while λ = 1 uses verifier scores.When the value model is inaccurate, lower λ has lower variance and higher bias, and higher λ the converse.Source: §2 · pp. 3–4
  • GRPO gives every token in a response the same advantage, computed from its terminal reward minus the group’s mean reward.Other RLVR methods may learn a value model but use it only as a baseline, retaining the λ = 1 estimator; this paper instead uses λ = 0 on some problems.Source: §2 · p. 4

3 Actor-Critic with Action Chunking

AC2 uses a replay buffer and a problem-specific readiness decision to choose between full rollouts and action chunks. It estimates chunk advantages from endpoint values, then updates the actor on generated tokens and trains the critic on group-mean targets.

3 Actor-Critic with Action Chunking

  • AC2 trains a policy and value model with shared parameters, using the critic for policy updates when its value estimates are accurate.The algorithm keeps complete prior rollouts in a replay buffer and samples both fresh problems and stored trajectories.Source: §3 · p. 4

3.1 Problem and policy sampling

  • For stored trajectories, AC2 samples random prefixes; unready problems receive full continuations, while ready problems usually receive chunks of at most b tokens.A fraction of ready problems are still fully rolled out to measure critic accuracy and update the critic; this auditing fraction is normally 1/4 and can be disabled.Source: §3.1 · p. 5

3.2 Advantage estimation

  • Each continuation’s endpoint value is its terminal reward if complete or the critic estimate otherwise; subtracting the group mean gives its advantage.Ready problems use the λ = 0 form of GAE and need no terminal reward, while unready problems use the λ = 1 form.Source: §3.2 · p. 6

3.3 Actor and critic update

  • The actor updates on newly generated tokens while replayed prefixes remain context, and the critic learns rounded group-mean targets.When a reference solution is available, the critic is trained on the same target with and without that solution in the prompt, at equal weight.Source: §3.3 · p. 6

4 Experiments

Across math-proof experiments, AC2 exceeds GRPO’s peak mean score with less decoding compute. The experiments and diagnostics compare AC2 with GRPO and variants while examining critic readiness, action-chunk size, audits, and replay choices.

4 Experiments

  • Experiments train Qwen3-4B-Thinking-2507 on FineProofs-RL and evaluate on IMO-ProofBench; Decoding FLOPs serve as a hardware-agnostic rollout-compute measure.GPU hours are not directly comparable across runs using different numbers of GPUs and communication overheads.Source: §4 · p. 7

4.1 Main Results

  • 0.79e20 FLOPs versus 1.99e20 FLOPs: AC2 first exceeds GRPO’s peak mean score of 18.50%; its peak mean score is 20.57% versus GRPO’s 18.5%.The authors report a 2.5× compute-efficiency gain over GRPO.Source: §4.1 · p. 7
  • Prefix GRPO falls well behind AC2 when measured by Decoding FLOPs, so replay-buffer state distribution alone does not explain AC2’s gains.Source: §4.1 · p. 7
  • Figure 2 tracks readiness and critic error; the readiness threshold is first crossed after step 7, and the ready-problem share reaches about 70% by step 200.With a reference solution in its prompt, the critic has lower MAE at every tested checkpoint.Source: §4.1 · pp. 7–8, Figure 2 · p. 8, §4.4 · p. 10

4.2 AC2 variants

  • AC2 without auditing peaks at 18.66%; the stale-buffer variant reaches 17.90%, while the one-action-chunk-per-ready-prefix variant performs similarly to AC2.The stale replay buffer is refreshed every ten steps; the single-chunk variant retains full-length groups for unready problems.Source: §4.2 · p. 8, Figure 3 · p. 9

4.3 Ablations that hurt performance

  • Correct-only replay reaches 17.35% at step 60 versus AC2’s 16.88%, but falls to 14.67% by step 110.The buffer keeps trajectories scoring at least six of seven training-judge points.Source: §4.3 · p. 9
  • At step 40, disabling local readiness yields 13.10% mean score versus 15.33% with readiness; the 2,000-token variant peaks at 16.62% and later falls to 13.57%.The 2,000-token run tracks the 10,000-token run through step 50, then stops improving; the authors report larger chunks are easier for the critic to judge.Source: §4.3 · p. 9

4.4 Critic Diagnostics

  • At step 80, group-mean values have MAE 0.065 versus 0.211 for prefix values; critic-based advantages have MAE 0.150 against Prefix GRPO advantages and Pearson correlation 0.388.The paper cautions that Prefix GRPO advantages should not be treated as ground truth because AC2 outperforms Prefix GRPO.Source: §4.4 · p. 10, §4.3 · p. 10, Figure 4 · p. 10

5 Related Work

The section contrasts terminal-reward group-relative methods with actor–critic approaches and explains how AC2 relates to both. It also situates AC2 among recent methods that use critics or intermediate values while retaining terminal rollouts.

5 Related Work

  • A3C estimates advantages from an n-step return and uses the critic at step n; AC2 similarly uses λ = 0 for a single action chunk without a terminal rollout.The paper notes the methods’ similarity while distinguishing their application settings.Source: §5 · pp. 10–11
  • Group-relative methods sample complete responses and give every token the same advantage, whereas actor–critic methods can assign denser credit across a trajectory.The section describes GRPO and descendants as using group-mean terminal rewards; actor–critic methods estimate advantages from shorter returns.Source: §5 · pp. 10–11
  • Recent critic-based RLVR methods still rely on terminal rewards, while Monte Carlo intermediate-value methods require additional terminal rollouts; AC2 uses λ = 0 instead.AC2 combines local critic readiness, optional correct-solution context, and action chunking; prior work explored these ingredients separately or with different value-estimation approaches.Source: §5 · p. 11

6 Conclusion

The conclusion frames AC2 as an example of more aggressive critic use in LLM RL, including updates that do not require terminal rewards. It identifies local readiness’s reliance on epoched data as a limitation and leaves adapting it to single-epoch training for future work.

6 Conclusion

  • Using a critic with λ = 0 removes the need to collect terminal rewards and, the authors argue, opens design space for LLM RL algorithms.AC2 is presented as one example; replay-buffer prefixes also make state resetting possible, with chunk beam search and MCTS suggested as possibilities.Source: §6 · pp. 11–12
  • The conclusion suggests combining critics with state resetting for other approaches, including beam search over action chunks or alpha-go-style MCTS.These are proposed possibilities, not evaluated results in the supplied conclusion passages.Source: §6 · pp. 11–12
  • Local readiness requires epoching training data because it uses complete rollouts to obtain ground-truth value estimates, which the authors say is unrealistic without data constraints.Adapting AC2 and local-readiness criteria to the single-epoch regime is left for future work.Source: §6 · p. 12

A Additional Experiments

The GRPO learning-rate sweep compares mean score against decoding FLOPs and training steps. The authors select 2 × 10−6 for their experiments based on its comparison with the other tested rates at step 100.

A.1 GRPO tuning

  • At step 100, 2 × 10−6 is better than 4 × 10−6 and comparable to 1 × 10−6, so the authors select it for experiments.Figure 5 shows mean score against decoding FLOPs and steps.Source: §A.1 · p. 17, Figure 5 · p. 17
  • Figure 5 presents the observed learning-rate sweep as mean score against decoding FLOPs and training steps.Source: Figure 5 · p. 17

A.2 Additional ablations

The additional ablations compare AC2 variants by training steps, including a 2,000-token chunk setting and a variant without local readiness. The chunk-size comparison changes prefix-cut spacing as well, so it does not isolate chunk length alone.

A.2 Additional ablations

  • Figure 6 shows the Figure 3 component comparisons against training steps; the starting-prefix baseline uses 16 separate prefixes with one continuation each, preserving 16 continuations per problem.Full groups instead retain 16 continuations per prefix.Source: Figure 6 · p. 17, §A.2 · p. 17
  • At step 50, AC2 w/ 2k chunks and AC2 reach mean scores of 16.03% and 16.41%, respectively; the 2k variant later peaks at 16.62% and ends at 13.57%.The comparison also changes prefix-cut spacing, so it does not isolate either change.Source: §A.2 · p. 18, Figure 7 · p. 18
  • Figure 8 plots sampled-problem ready fractions for AC2 w/o Group & Audit from the base model and a variant that also removes local readiness from AC2 at step 20.Source: Figure 8 · p. 18

A.3 Additional evaluation metrics

On best-of-16 validation score, AC2 remains more compute efficient than GRPO, although both plateau at roughly the same score.

A.3 Additional evaluation metrics

  • AC2 remains more compute efficient than GRPO on best-of-16 score, while the two methods plateau at roughly the same score.The score is estimated by bootstrapping 16 sampled validation trajectories and taking the maximum in each sample.Source: §A.3 · p. 18
  • Figure 9 presents best-of-16 scores for AC2 and GRPO.Source: Figure 9 · p. 19

A.4 Additional value function diagnostics

These diagnostics compare critic-derived continuation values with terminal rewards across checkpoints and groupings. Checkpoints use different sampled problems, so their differences do not measure improvement on a fixed set.

A.4 Additional value function diagnostics

  • At each checkpoint, the diagnostics sample 256 ready problems, cut latest trajectories into prefixes, and generate 16 continuations per prefix.Each continuation’s first 10k tokens, or fewer if it ends sooner, form its action chunk; completed continuations use terminal reward, otherwise the critic supplies the value.Source: §A.4 · p. 19
  • The comparisons cover group-mean values versus rewards, critic estimates on eligible continuations, and critic-based versus GRPO advantages.Additional views compare prefix critic values with terminal rewards and inspect group means and individual continuation estimates at different steps.Source: §A.4 · pp. 19–20, Figure 10 · p. 20
  • Checkpoint comparisons do not establish improvement on a fixed problem set because the sampled problems differ between checkpoints.Groups with f = 0 have matching value and reward means by construction; only groups with f > 0 test the critic.Source: §A.4 · p. 19
  • The group-mean value-versus-reward comparison is shown at steps 40 and 57; the figure’s color encodes the share of continuations valued by the critic.Source: Figure 12 · p. 21
  • The diagnostics also compare individual chunk-end critic values with terminal rewards at steps 40, 57, 80, and 160.Source: Figure 17 · p. 24, Figure 18 · p. 24

B Decoding FLOPs accounting and GPU hours

The paper estimates decoding FLOPs from per-request prompt and generated-token lengths, using exact joint statistics when available and mean-based estimates otherwise. In the main AC2 run, decoding FLOPs are strongly associated with GPU hours.

B Decoding FLOPs accounting and GPU hours

  • The per-request decoding FLOP calculation uses prompt-plus-prefix length and newly generated response length, with architecture-specific coefficients for forward and attention computations.The accounting excludes prefill, validation, value-function and judge calls, training computation, elementwise operations, and communication.Source: B Decoding FLOPs accounting and GPU… · p. 25
  • Cumulative decoding FLOPs sum costs across actor-generated tokens, including tokens used to refill the replay buffer.When paired prompt-plus-prefix and generated lengths are unavailable, the estimate substitutes empirical mean lengths into the cost formula.Source: B Decoding FLOPs accounting and GPU… · p. 25
  • Mean-based cost estimates can over- or underestimate true costs, and the GRPO curves use these estimates while AC2 uses exact records.For available records, the estimate is 4.01% below exact cost for the selected GRPO run and 5.51% below for main AC2; the authors note plotted GRPO cost is, if anything, slightly low.Source: B Decoding FLOPs accounting and GPU… · p. 25
  • In the main AC2 run, per-step decoding FLOPs correlate with GPU hours at r = 0.874.The fitted relation is bH = 6.11 + 27.41 (D/1018), based on 196 complete steps.Source: B Decoding FLOPs accounting and GPU… · p. 26, Figure 19 · p. 26

C Training and evaluation details

The training details specify replay-buffer sampling, reference-solution use, readiness-based rollout choices, and optimization and evaluation settings. Ready problems may receive chunks rather than full rollouts, while unready problems and audited ready problems receive full rollouts.

C Training and evaluation details

  • The replay buffer stores up to 256 trajectories, and prefixes are cut at multiples of 10,000 tokens up to 90% of response length.The total response budget includes the replayed prefix and newly generated tokens.Source: C Training and evaluation details · pp. 27–28
  • A reference bank stores one proof per problem when a rollout earns at least six of seven points; existing references are not replaced.Readiness uses the bank before additions at the current step, though critic training can use newly added references immediately.Source: C Training and evaluation details · p. 28
  • Critic targets require at least eight valid rewards in a continuation group, and targets are rounded to the nearest tenth from 0 to 1.A reference-paired training example is trained under both prompts at weight one-half each; examples without a reference use the plain prompt at weight 1.Source: C Training and evaluation details · p. 28
  • At each step, one quarter of sampled ready problems are audited with full rollouts; other ready problems receive action chunks, while unready problems always receive full rollouts.A bug activated the reference-proof readiness condition only at step 42 in the main run; the authors expect this to have hurt reported AC2 performance versus activating it from step 1.Source: C Training and evaluation details · p. 29
  • Invalid endpoint critic values receive zero actor-loss weight, while completed responses, budget-exhausting responses, or responses containing a complete proof are judge-scored.Source: C Training and evaluation details · p. 29
  • Each step performs an actor update, a critic update, and a second actor update; evaluation samples 16 responses per problem every ten training steps.The main evaluation response limit is 50,000 tokens, and equal training steps need not represent equal response counts or decoding work.Source: C Training and evaluation details · p. 29
Loading 2609.39247v2…