PACT: From Credit Assignment to Critic Alignment

AI summary

Overview

Research area: Reinforcement learning (RL) for large language model (LLM) post-training, specifically token-level credit assignment and actor-critic training for long-horizon, agentic tasks.

Technical level: Advanced. The core of the paper is a measure-theoretic representation theorem (filtration-adapted conditional expectations, martingale difference sequences) followed by a practical training algorithm.

Scope: The paper defines token-level credit through three regularity conditions, proves the representation is unique, uses it to reinterpret existing algorithms (On-Policy Distillation, RLOO, GAE), and proposes Policy Aligned Critic Training (PACT), which is evaluated on four mathematical reasoning benchmarks and SWE-bench Verified.

What This Paper Is About

In LLM post-training, a model generates a long sequence of tokens (often interacting with tools or environments), but the training signal is usually a single scalar reward revealed only at the end. How that final outcome should be attributed to individual tokens—the credit assignment problem—has no generally accepted mathematical definition, so different algorithms operationalize credit in different ways. The paper sets out to give credit a precise characterization and then uses that characterization to diagnose and fix weaknesses in existing actor-critic training.

Key Contributions

1. A unique representation theorem for token-level credit. The authors formulate three regularity conditions—Completeness, Prefix Consistency, and Neutrality—and prove that a token-level credit assignment satisfying them exists and is unique up to almost-sure equality, given by consecutive differences of conditional reward predictions, C_i = V_i − V_{i−1} = E[R | F_i] − E[R | F_{i−1}], forming a martingale difference sequence. Appendix C.2.3 is stated to show all three conditions are necessary, since omitting any one admits alternative assignments.

2. A unified reinterpretation of existing algorithms. Under the unique credit, an ideal teacher in On-Policy Distillation acts as an implicit critic (its expected policy gradient is proportional to the credit-induced gradient, up to a factor 1/β); the response-level RLOO leave-one-out signal matches the expected policy-gradient contribution of token-level credit despite its coarser granularity; and approximate credit sparsity under bounded outcome rewards explains why intermediate critic errors can dominate in GAE with λ < 1.

3. Policy Aligned Critic Training (PACT). Motivated by the analysis, PACT adopts an Actor-then-Critic update order so that importance sampling correction can be applied to critic training, aligning the critic with the updated policy. It also replaces Mean Squared Error with Binary Cross Entropy for value estimation and uses λ = 1 in GAE.

4. Empirical validation on reasoning and coding. PACT reports 72.87% average accuracy across four mathematical reasoning benchmarks and a 67.4% pass rate on SWE-bench Verified.

Main Findings

• Uniqueness of credit: Completeness, Prefix Consistency, and Neutrality together uniquely determine token-level credit as C_i = V_i − V_{i−1}, where V_i = E[R | F_i]; coarser assignments such as turn-level credit are recovered by aggregating tokens over consecutive segments (Corollary 1).

• The OPD teacher is an implicit critic: With the ideal teacher defined as the KL-regularized maximizer of Q_t^π(a), the OPD update direction satisfies G_t^OPD(q_t⋆) = (1/β) E_π[Z_t C_t^π | F_{t−1}] almost surely—i.e., the teacher's signal is proportional to the credit-induced gradient.

• RLOO is gradient-equivalent but not efficiency-equivalent: The leave-one-out baseline R̄_{−i} (built from G ≥ 2 trajectories) yields the same expected token-level policy gradient as the unique credit, but the equivalence holds only in expectation; the response-level baseline retains randomness from finite outcomes, which can be large relative to local credit in long-horizon settings.

• Approximate credit sparsity: For R ∈ [0,1], the expected sum of squared credits equals the conditional reward variance, which is at most 1/4. Consequently the expected number of credits exceeding any ε > 0 is at most 1/(4ε²), independent of maximum response length.

• Why λ = 1 helps in GAE: Decomposing the estimated advantage shows that intermediate critic errors enter through a (1 − λ) weighted term when λ < 1, while at λ = 1 the advantage reduces to Â_t^1 = R − V̂_{t−1}, eliminating intermediate value errors and leaving only the prefix value error. PACT therefore uses λ = 1.

• Mathematical reasoning results (Avg@16 accuracy): Base model 41.01% average; GRPO (ε_high = 0.28) 64.07%; PPO (λ = 0.95) 26.31%; PPO (λ = 1.0) 59.71%; SAO 51.14%; PACT without importance sampling 67.74%; PACT 72.87%. PACT outperforms GRPO, PPO (λ = 1.0), and SAO by 8.80, 13.16, and 21.73 percentage points respectively, and achieves the highest accuracy on all four benchmarks: AIME 2025 83.12, AIME 2026 85.21, BeyondAIME 51.69, HMMT Nov. 2025 71.46.

• Policy collapse at λ = 0.95: PPO with λ = 0.95 undergoes policy collapse during training, consistent with the paper's analysis of intermediate critic errors in GAE.

• SWE-bench Verified results (pass@1): Base model 60.8%, GRPO 65.4%, PPO (λ = 1.0) 65.0%, SAO 63.6%, PACT 67.4%. PACT outperforms GRPO, PPO (λ = 1.0), and SAO by 2.0, 2.4, and 3.8 percentage points respectively.

• Critic objective ablation: With a fixed rollout policy, critics trained from the same initialization on the same on-policy rollout data using BCE achieve lower BCE and MSE losses and greater value separation (Δ±, the mean predicted value on positive samples minus that on negative samples) than MSE-trained critics, across both model scales.

• Importance sampling ablation: Removing the critic-target importance sampling correction while keeping the Actor-then-Critic order, BCE objective, and actor-side settings lowers average accuracy from 72.87% to 67.74%—a 5.13 percentage-point gain attributed to the correction—with improvements on all four benchmarks and more stable training reward curves (Appendix E.3).

Methodology in Plain English

The authors begin by writing down what any sensible notion of per-token credit should satisfy. First, the credits summed over a trajectory should exactly equal how much better the final reward was than what was expected before generation started. Second, once a prefix has been generated, the credit accumulated so far should not change based on what happens later. Third, the credit assigned to a token should not be systematically biased given everything known beforehand—its conditional expectation is zero. They then prove these three requirements leave exactly one possible answer: the credit for a token is the change in the model's conditional prediction of the final reward caused by that token.

With this fixed answer in hand, they re-derive what existing algorithms are actually doing. They show that an idealized distillation teacher plays the same role an explicit value network would; that a leave-one-out baseline computed across several sampled responses gives the same gradient in expectation even though it works at the response level; and that because most token credits are tiny when rewards are bounded in [0,1], any error in the value estimate can swamp the true signal—especially when GAE mixes errors across positions with λ < 1.

The method that follows is straightforward. Within each training iteration, the critic first supplies values for the actor update. After the actor has been updated, the same rollout batch is run through the new actor to compute importance ratios between the pre-update and post-update policies, and the critic is then trained on importance-corrected targets so that the iteration ends with the critic matching the actor that will actually be used next. The critic is parameterized as a sigmoid over a logit and trained with binary cross-entropy instead of squared error (for rewards in [0,1], both objectives have the same optimum). To control the variance of the exact continuation ratio in long responses, the implementation uses the detached current-token importance ratio and masks out token-level critic losses whose ratios fall outside [0, 6].

Experiments use the Dressage agentic RL framework built on slime. Mathematical reasoning: Qwen3.5-4B trained on a 3,200-problem subset of DAPO-Math-17k with the OpenCode agentic harness, using final-answer correctness as the reward, evaluated with Avg@16 accuracy on AIME 2025, AIME 2026, BeyondAIME, and HMMT Nov. 2025. Coding: Qwen3.6-35B-A3B trained on OpenSWE using a Codex agent through Harbor, with terminal rewards from the task verifier, evaluated by pass@1 on SWE-bench Verified. Each rollout round produces 512 trajectories; GRPO samples 8 trajectories for each of 64 prompts; the optimization batch size is 128, yielding four minibatches per rollout round. SAO uses DIS with importance ratio ranges of [0.7, 6.0] for mathematics and [0.6, 3.0] for coding; PACT uses the same actor-side DIS range as SAO for mathematics and PPO clipping for coding. Context window is 128k tokens and maximum generation length per interaction turn is 64k tokens.

Why This Matters

Impact on research. The paper turns an informal notion—credit—into a characterized mathematical object, then shows that several apparently different algorithms (distillation with a teacher, group-relative baselines, critic-based advantage estimation) are connected to that object. This gives a common reference point for comparing methods and isolates a specific mechanistic explanation (intermediate critic error under λ < 1) for a documented failure mode: policy collapse of PPO at λ = 0.95 in these experiments.

Real-world applications (the paper evaluates two domains; the others are natural analogues that the paper does not itself test):

• Agentic coding assistants that must decide which edits, tool calls, or commands in a long session deserve credit for a task passing or failing, as evaluated on SWE-bench Verified.
• Mathematical reasoning and competition-style problem solving with tool use, as evaluated on AIME 2025, AIME 2026, BeyondAIME, and HMMT Nov. 2025.
• Multi-turn tool-using agents where a single reward arrives after many environment observations—the general setting the notation covers, since observations may be tool responses, environment transitions, or user feedback.
• Long-horizon RL post-training pipelines that would otherwise need to train a value model separately, where the analysis of λ and critic error offers guidance on configuration choices.

Industry relevance. The method adds one extra forward pass over an existing rollout batch and requires no additional rollout generation, which keeps the cost profile close to standard PPO-style training. The GitHub repository is listed as https://github.com/AllSpark-Research/PACT, and the training stack (Dressage on slime, OpenCode, Codex/Harbor, Qwen3.5-4B and Qwen3.6-35B-A3B) is described in concrete detail, which lowers the barrier to reproduction.

Future Directions

• Beyond bounded rewards. Theorem 4's sparsity guarantees are stated under R ∈ [0,1], obtained by normalizing a bounded reward with a positive affine transformation. Whether analogous guarantees hold for unbounded or non-normalizable reward signals is not established in the paper.
• Variance control for the importance ratio. The paper notes that the exact continuation ratio I_t can have high variance for long responses, and the implementation substitutes a detached current-token ratio with masking outside [0, 6]. Better estimators for policy-synchronized critic targets remain an open practical question.
• Critic synchronization across more than one update. The PACT scheme addresses the case where the critic is one policy update behind; whether the Actor-then-Critic dependency should be extended over longer horizons, or combined with asynchronous rollout systems described in the related work, is not resolved.
• Coarser granularities and other domains. Corollary 1 shows turn-level and other segment-level credit follow from aggregating token-level credit over consecutive segments, but the experiments cover only two settings (agentic mathematical reasoning and agentic coding). Applying the characterization to other agentic environments is left open.

Target Audience

Researchers and engineers working on RL post-training for LLMs who are comfortable with probability theory and policy gradient methods—particularly those building critic-based (rather than purely critic-free) training pipelines for long-horizon agentic tasks. Readers interested in the theory of credit assignment will find the representation theorem and its consequences the main draw; practitioners will find the PACT update order, the BCE critic objective, and the λ = 1 recommendation the most directly transferable elements. A reader looking for a purely empirical benchmark report would need to go beyond the two evaluation suites included here.

Authors’ abstract

Reinforcement learning has become a central component of large language model (LLM) post-training, yet token-level credit lacks a generally accepted mathematical definition, leaving its relationship to commonly used training signals unclear. We formulate three regularity conditions, namely Completeness, Prefix Consistency, and Neutrality, and prove that they uniquely determine token-level credit. This characterization provides a unified basis for explaining phenomena across existing algorithms and guides the development of an improved actor-critic training procedure. Through this lens, an ideal teacher in On-Policy Distillation (OPD) acts as an implicit critic, yielding an expected policy gradient proportional to that induced by token-level credit. Response-level REINFORCE Leave-One-Out (RLOO) signals match the expected policy-gradient contribution of token-level credit despite their coarser granularity. We further establish approximate credit sparsity under bounded outcome rewards and show how intermediate critic errors in Generalized Advantage Estimation (GAE) can become comparable to the underlying credit. These motivate Policy Aligned Critic Training (PACT), which adopts an Actor-then-Critic update order to apply importance sampling correction to critic training and better align the critic with the updated policy. In agentic mathematical reasoning, PACT achieves 72.87% average accuracy across four benchmarks, outperforming GRPO and PPO by 8.80 and 13.16 percentage points, respectively. On SWE-bench Verified, PACT achieves a pass rate of 67.4%, outperforming PPO, GRPO, and SAO by 2.4, 2.0, and 3.8 percentage points, respectively.