Technical Summary: CATPO (Critique-Augmented Tree Policy Optimization)

Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has become a dominant paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs). While tree-based methods like TREERPO improve upon flat trajectory sampling by generating dense, step-level reward signals through tree-structured rollouts, they suffer from a critical inefficiency: computational waste due to uninformative trees.

Current tree-based approaches treat all sampled trees equally during training. However, many trees provide little to no useful gradient signal:

• Dead-correct trees: All leaves succeed. The model has already mastered the problem, offering no contrastive signal for improvement.
• Dead-wrong trees: All leaves fail. Without any positive branches, group-relative advantages collapse to zero, rendering the rollout useless for gradient updates.
• Stale trees: The policy's confidence already perfectly tracks the reward distribution, yielding diminishing returns.

Existing methods like DAPO and GRESO address similar issues in flat rollouts by filtering prompts, but they operate at the prompt level. They cannot diagnose or repair the internal structure of a tree rollout, where specific branches may be surprising or erroneous while others are not.

Methodology

CATPO (Critique-Augmented Tree Policy Optimization) addresses this waste through a four-phase pipeline that diagnoses, repairs, and re-weights tree rollouts at the tree level.

1. Tree Informativeness Score ($F(T)$)

CATPO introduces a metric to quantify the utility of a tree $T$ for policy updates. This score is computed at zero additional cost using data already generated during sampling. It combines two signals:

• Leaf Outcome Diversity: Measured by the Bernoulli variance of the leaf pass rate ($\hat{p}$), calculated as $\hat{p}(1-\hat{p})$. This term is maximized when outcomes are mixed (some correct, some incorrect) and zero when all leaves agree.
• Policy-Reward Decorrelation: Measured by $1 - \rho^2_{\pi,r}$, where $\rho_{\pi,r}$ is the Pearson correlation between node log-probabilities and propagated rewards. This captures the residual uncertainty where the policy's confidence does not linearly explain the reward structure.

The final score is $F(T) = \hat{p}(1-\hat{p}) \cdot (1 - \rho^2_{\pi,r})$. Trees are classified into four regimes based on thresholds:

• Dead-correct/Dead-wrong: Low diversity (near-unanimous outcomes).
• Stale: Moderate signal.
• Informative: High diversity and decorrelation.

2. Critique-Guided Healing

For dead-wrong trees (where all branches fail), CATPO does not discard the rollout. Instead, it attempts to "heal" the tree:

1. Failure Identification: A breadth-first search locates the shallowest node $v^*$ where all children have failed (reward < $\epsilon$).
2. Critique Generation: The policy model itself is prompted to generate a natural-language critique of the error at node $v^*$, identifying the specific logical or mathematical mistake.
3. Grafting: The model is prompted again with the problem, the partial solution up to $v^*$, and the generated critique to produce $k$ refined continuations.
4. Re-propagation: These new branches are grafted onto the tree at $v^*$, re-scored, and rewards are propagated bottom-up. This converts a "dead" tree into one with positive reward diversity.

3. Informativeness-Weighted Loss

The final training objective scales the contribution of each tree to the policy gradient by its normalized informativeness score $w(T) = F(T) / \bar{F}$.

• Informative trees receive full weight, driving parameter updates.
• Dead-correct and stale trees are downweighted.
• Healed trees are treated as informative if the healing process succeeds.

The objective function uses a clipped surrogate loss (similar to PPO/TREERPO) where the advantages are modulated by $w(T)$, concentrating updates on the most useful data while preserving overall gradient magnitude.

Key Contributions

1. Tree-Level Diagnosis: CATPO is the first method to explicitly diagnose tree informativeness and intervene on uninformative trees within the RL training loop, moving beyond prompt-level filtering.
2. Zero-Cost Scoring: The informativeness score $F(T)$ leverages existing tree sampling statistics (leaf outcomes and log-probs) without requiring auxiliary models or extra forward passes.
3. Critique-Guided Healing: A novel mechanism that recovers training signal from "dead-wrong" trees by using the policy model to critique its own errors and graft corrected branches, effectively turning wasted compute into useful training data.
4. Weighted Optimization: An informativeness-weighted loss that dynamically prioritizes high-value trees, improving sample efficiency.

Experimental Results

Experiments were conducted on Qwen2.5-Math-1.5B trained on the MATH dataset and evaluated across four benchmarks: AIME24, MATH-500, OlympiadBench, and MinervaMath.

• Performance: CATPO achieved a 37.5% Macro Accuracy, outperforming the baseline TREERPO (35.6%) by 1.9% and the flat-trajectory GRPO (32.7%) by 4.8%.
• Harder Tasks: The improvements were most pronounced on difficult benchmarks (AIME24 and OlympiadBench), where the baseline produces a higher proportion of dead-wrong trees, giving CATPO more opportunities to recover signal.
• Generalization: Significant gains were observed on MinervaMath (outside the training distribution), suggesting the method helps learn transferable reasoning strategies rather than overfitting.
• Hypothesis Validation:
  • H1: The informativeness score $F(T)$ showed a strong positive correlation (Spearman $\rho = 0.62$) with per-tree gradient norms, validating it as a reliable proxy for training utility.
  • H2: Critique-guided healing successfully improved the informativeness of 88.1% of dead-wrong trees by introducing at least one correct branch.

Significance and Claims

The paper claims that CATPO offers a more efficient and effective approach to RLVR for LLM reasoning by directly addressing the "waste" inherent in tree-based sampling. By distinguishing between informative and uninformative trees and actively repairing the latter, CATPO concentrates computational resources on the data that drives learning. The authors emphasize that this approach requires no separate reward models or value networks, maintaining the simplicity of the GRPO/TREERPO paradigm while significantly boosting performance on complex reasoning tasks. The method is presented as a complementary advancement to existing tree-based RL, specifically targeting the internal structure of rollouts that previous methods ignored.