Multiscale Policy Optimization for Long-Horizon LLM Agents

AI summary

Overview

Research area: Reinforcement learning (RL) post-training for large language model (LLM) agents that act over many turns in an environment — search, tool use, and computer/automation tasks.

Technical level: Advanced. The paper pairs a formal policy-improvement analysis (a turn-level improvement bound plus a proof that token-level surrogates pay an extra within-turn distribution-shift penalty) with a new policy-gradient objective and a large empirical study.

Scope (1 sentence): The paper argues that credit assignment and importance weighting for LLM agents should happen at the scale of complete assistant turns — while still distinguishing tokens inside a turn — and introduces Multiscale Policy Optimization (MPO) to do both, validated on search, tool-use, and automation benchmarks with 4B and 30B-A3B models.

Note: the supplied text ends mid-sentence in Appendix A.1, so appendices beyond the start of that proof (including Appendix B and parts of C and D referenced in the body) are not available for summary.

What This Paper Is About

An LLM agent produces whole assistant turns (a search query, a tool call, a final answer) but writes each turn one token at a time, and typically receives only a single reward at the end of the trajectory. Existing RL methods pick one scale: group-relative methods like GRPO give one shared advantage per sequence, while critic-based token-level methods like PPO treat every token as an action without separating turn-level interaction from within-turn generation. The goal is a policy-optimization method whose credit assignment and probability correction both match the turn structure of agentic interaction without discarding token-level detail.

Key Contributions

1. Turn-level and token-level theory of policy improvement. The authors establish a finite-horizon turn-level improvement bound (Theorem 1) treating complete assistant outputs as macro actions, then show that a token-level surrogate under the same KL constraints incurs an additional within-turn distribution-shift penalty (Theorem 2).
2. An empirical account of critic structure in agent rollouts. Across 2,048 trajectories and 16,323 turns from eight checkpoint batches of one search-agent run, differences between turn means account for a median of 89.5% of the critic's token-weighted within-trajectory value variance, and between-turn variance exceeds within-turn variance in 93.8% of trajectories.
3. Multiscale Advantage Estimation (MAE). Turn-level GAE propagates credit across interactions, and a centered local correction (derived as a constrained least-squares projection with a single critic) refines token credit while preserving each turn's mean advantage.
4. Turn-level importance correction inside the gradient. A stop-gradient construction gives every token the complete-turn importance ratio in the forward pass and clipping, while each token's own MAE advantage and score gradient remain token-specific; the result is the clipped MPO objective.
5. Consistent empirical gains over GRPO, GSPO, PPO, VAPO, and SAO across nine metrics at two model scales, plus ablations isolating each scale choice.

Main Findings

• Critic values are turn-structured. The variance decomposition (median 89.5% between-turn; between-turn variance larger in 93.8% of trajectories) motivates propagating long-range credit at the turn scale rather than the token scale.
• MPO leads every reported metric at both scales. Table 1 shows MPO highest on all nine metrics for Qwen3-4B-Thinking and Qwen3-30B-A3B-Instruct. Against the strongest non-MPO baseline per column, MPO improves BrowseComp-Plus, GAIA, and HLE pass@1 by 4.33, 3.85, and 0.66 points on 4B, and by 3.59, 3.80, and 0.12 points on 30A3B.
• Gains hold at pass@3 and on multi-trial success. GAIA pass@3 improves by 5.11 and 5.90 points at the two scales, and Claw-Eval pass^3 rises by 0.95 and 1.12 points.
• Turn-level importance ratios matter, but MAE carries gains alone. Replacing complete-turn ratios with token-level ratios ("MPO w/ token IS") reduces every reported score — Search pass@1 drops of 2.36, 3.12, and 1.03 points on BrowseComp-Plus, GAIA, and HLE for 4B, and 2.05, 2.19, and 0.58 for 30A3B — yet this variant still beats PPO on all six Search metrics at both scales.
• Benefits persist under larger sampling budgets. On 4B, MPO leads GRPO and PPO at every evaluated k in {1, 2, 4, 8, 16, 32}; its pass@32 reaches 75.90 on BrowseComp-Plus, 69.90 on GAIA, and 34.00 on HLE, exceeding the stronger of GRPO and PPO by 3.61, 5.82, and 5.80 points.
• Both scales are needed. With β = 0 (turn broadcasting only), the method still beats PPO on all six Search metrics, with pass@1 gains of 0.95, 2.49, and 0.66 points, but removing local allocation costs 6.28, 5.27, and 2.04 pass@1 points relative to β = 1.00.
• Coefficients should differ by scale. The best setting is (λ, λ_intra, β) = (0.80, 0.95, 1.00). Raising λ to 0.85 reduces pass@1 by 1.72, 2.21, and 1.05 points; lowering λ to 0.60 also hurts. Reducing λ_intra from 0.95 to 0.90 lowers pass@1 by 1.83, 1.55, and 1.77 points, and λ_intra = 1 lowers all six scores. The nominal trace length 1/(1−λ) is five turns at λ = 0.80 but twenty tokens at λ_intra = 0.95.
• Critic fit improves under MPO. With λ_critic = 1 for both PPO and MPO, MPO shows higher explained variance on average during late training at both scales; the authors flag this as a possible interaction effect rather than a causal claim, since explained variance is measured on each policy's own rollouts.
• Entropy alone is not collapse. PPO's entropy drops sharply during 30A3B training while validation performance keeps improving, and lower 4B MPO entropy coexists with strong pass@k, which the authors treat as direct evidence about exploration.

Methodology in Plain English

The authors start by formalizing an agent trajectory as alternating assistant outputs and environment observations, with a single reward at the very end (only the final assistant token receives it; all earlier tokens get zero). Because an assistant turn is one action from the environment's point of view, they borrow the classical trust-region argument and redo it with turns as the unit of action, obtaining a bound that combines a turn-level advantage with the ratio of complete-turn probabilities. Comparing that bound to a token-level surrogate shows the token version is penalized by how much the token distributions shift inside a turn — a penalty that grows with maximum turn length.

The method follows from this. In a first stage, temporal-difference residuals are summed within each turn so internal value predictions cancel, leaving only boundary values and the terminal reward; standard GAE is then applied along the turn axis, so credit decays once per interaction rather than once per token. In a second stage, a local trace is computed inside each turn, and the turn advantage plus a scaled, mean-removed version of that local trace gives per-token advantages — a solution to a constrained least-squares projection that keeps the turn mean fixed while preserving relative token differences. A single critic serves both stages, and with λ_critic = 1 it regresses the terminal return directly, following VAPO's decoupled GAE.

For the update, the product of token probability ratios gives the complete-turn ratio. Multiplying all advantages by that same ratio would cancel the centered local correction, so a stop-gradient is applied to all other token ratios: every token sees the full turn ratio in the forward pass and during clipping, but the gradient carries the turn ratio times that token's own log-probability and its own MAE advantage. The resulting clipped surrogate is then optimized.

Experiments train Qwen3-4B-Thinking and Qwen3-30B-A3B-Instruct separately in each setting (300 steps for Search, 300 for Claw-Eval, 200 for Automation Bench). The best checkpoint per run is evaluated offline five times and averaged; trained Search models estimate pass@1 and pass@3 from 32 independent rollouts per task, Claw-Eval uses three independent trials (pass@3 and pass^3), and Automation Bench uses pass@1 with end-state assertions. All reported score differences are in percentage points.

Why This Matters

Impact on research. The paper reframes a practical design question — which scale to assign credit and correct probabilities at — as an analyzable one, and gives a concrete penalty term explaining why token-level surrogates can be mismatched to agentic interaction. It also supplies a variance-decomposition measurement of critic behavior that other agentic-RL papers can replicate or challenge.

Real-world applications:

• Deep-research and web-search agents that issue many queries before answering (BrowseComp-Plus, GAIA, HLE).
• Tool-calling assistants that chain APIs, retrieval, and code execution across turns.
• Computer-use and workflow automation agents judged by end-state assertions (Automation Bench).
• Multi-attempt agentic deployments where a fixed checkpoint is sampled many times (Claw-Eval's three-trial protocol, pass@k up to 32).

Industry relevance. Agentic post-training is expensive because rollouts span many turns and long token sequences; a method that improves pass@1 and pass@3 at both 4B and 30B-A3B scales, and keeps improving as sampling budget grows, is directly relevant to teams deciding how to spend RL compute on agent products.

Future Directions

• Adaptive scale coefficients. λ and λ_intra were tuned one at a time on 4B Search with a single run per setting; whether the (0.80, 0.95) pair and β = 1.00 transfer to other horizons, tools, and model families is untested.
• Understanding the critic–actor interaction. The higher explained variance under MPO is explicitly not isolated as causal; separating changes in visited histories from genuine critic-accuracy gains is left open.
• Combining with other update regimes. The paper cites recent work on stable updates for asynchronous agent training but does not evaluate MPO in that setting.
• Behavior under estimated advantages. The theoretical results are derived for exact advantages, with an extension to estimated turn advantages placed in an appendix not included in the supplied text; empirical behavior under a poorly fit critic remains a practical question.
• Longer horizons and larger budgets. Test-time scaling is reported up to k = 32 and training dynamics to step 600 (4B) and step 300 (30A3B); behavior beyond these ranges is not reported.

Target Audience

RL and post-training researchers working on LLM agents; engineers building multi-turn tool-use, search, or computer-use agents who need to choose between group-relative and critic-based objectives; and theoretically inclined readers interested in how trust-region-style policy-improvement bounds change when the action unit is a complete generated turn rather than a token. Readers should be comfortable with policy gradients, GAE, importance sampling, and KL/TV-based trust-region arguments.

Authors’ abstract

Long-horizon LLM agents interact with their environments through complete assistant turns while generating each action token by token. We observe that critic values vary primarily across turns, motivating optimization anchored at turns with finer credit allocation among tokens. We propose Multiscale Policy Optimization (MPO), combining Multiscale Advantage Estimation (MAE) and turn-level importance sampling to align credit assignment and policy updates at the turn scale. Extensive experiments across search, tool-use, and automation tasks with 4B and 30B-A3B models demonstrate consistent gains over strong RL baselines.