Multiscale Policy Optimization for Long-Horizon LLM Agents

Summary (Overview)

• Core Problem: Long-horizon LLM agents interact with environments through complete assistant turns (e.g., search queries, tool calls) while generating each action token-by-token. Existing RL methods either use group-relative rewards that ignore turn/token distinctions (GRPO, DAPO, GSPO) or token-level critics that don't explicitly separate within-turn generation from turn-level interaction (PPO, VAPO, SAO).

• Key Observation: Critic values vary primarily across turns rather than within turns—between-turn variance accounts for a median of 89.5% of the critic's token-weighted within-trajectory value variance across 2,048 trajectories and 16,323 turns.

• Main Contribution: Proposes Multiscale Policy Optimization (MPO), combining Multiscale Advantage Estimation (MAE) with turn-level importance sampling to align credit assignment and policy updates at the turn scale while preserving token-specific credit.

• Theoretical Foundation: Derives a turn-level policy improvement bound (Theorem 1) and shows token-level surrogates incur an additional within-turn distribution-shift penalty under the same KL constraints (Theorem 2).

• Results: Consistent performance gains over strong RL baselines (GRPO, GSPO, PPO, VAPO, SAO) across search, tool-use, and automation tasks with both 4B and 30B-A3B models, achieving the highest scores on all nine metrics at both model scales.

Introduction and Theoretical Foundation

Background and Motivation

LLM agents solve tasks by alternating between assistant outputs (turns) and environment observations. Each turn is generated token-by-token autoregressively. Learning from a terminal task reward involves two scale choices:

1. Credit assignment: Connecting final outcomes to interactions across the trajectory while distinguishing tokens within each interaction
2. Policy updates: Accounting for changes in complete-output probabilities while retaining local token distinctions

Formal Setup

The interaction trajectory is defined as:

$$\tau = (o_0, a_0, o_1, a_1, \ldots, o_T, a_T).\tag{1}$$

where $o_0$ is the task prompt, $a_t$ is an assistant output, and $o_{t+1}$ is environment feedback. The policy factorizes autoregressively:

$$\pi_\theta(a_t \mid h_t) = \prod_{i=1}^{L_t} \pi_\theta(a_{t,i} \mid h_{t,i}).\tag{5}$$

The agent receives a single terminal reward $R(\tau)$ with no intermediate rewards:

$$r_{t,i} := \begin{cases} R(\tau), & t = T, i = L_T, \\ 0, & \text{otherwise}. \end{cases}\tag{4}$$

Key Theoretical Results

Lemma 1 (Objective consistency): Under lossless token refinement, identical reward/stopping rules, and $\gamma = 1$, the macro and token descriptions have the same task return, with:

$$\nabla_\theta \log \pi_\theta(a_t \mid h_t) = \sum_{i=1}^{L_t} \nabla_\theta \log \pi_\theta(a_{t,i} \mid h_{t,i}).\tag{7}$$

Theorem 1 (Monotonic improvement at turn scale): With $\alpha_M = \sup_{t,h} D_{TV}(\pi'(\cdot|h), \pi(\cdot|h))$:

$$\mathcal{J}(\pi') - \mathcal{J}(\pi) \geq \mathcal{S}_M - 2A_{\max}T(T+1)\alpha_M^2.\tag{10}$$

Theorem 2 (Additional cost of token surrogates): The difference between macro and token surrogates is bounded by:

$$|\mathcal{S}_M - \mathcal{S}_T| \leq 2(T+1)A_{\max}^{\text{tok}}L_{\max}(L_{\max}-1)\alpha_T^2.\tag{11}$$

Takeaway: Under the same KL constraints, token-surrogate improvement bounds incur an additional within-turn distribution-shift penalty relative to turn-surrogate bounds.

Methodology

Multiscale Advantage Estimation (MAE)

MAE operates in two stages, both reusing the critic $V_{t,i} = V(h_{t,i})$:

Stage 1: Turn-Level Credit Propagation

Summing TD residuals within a turn cancels internal value predictions:

$$\delta_t^{\text{turn}} = \sum_{i=1}^{L_t} \delta_{t,i} = \mathbf{1}\{t=T\}R(\tau) + V_{t+1,1} - V_{t,1}.\tag{16}$$

Applying GAE on the macro time axis:

$$A_t^{\text{turn}} = \delta_t^{\text{turn}} + \lambda A_{t+1}^{\text{turn}} = \sum_{\ell=0}^{T-t} \lambda^\ell \delta_{t+\ell}^{\text{turn}}.\tag{17}$$

Stage 2: Within-Turn Credit Allocation

A local trace restricted to the current turn:

$$A_{t,i}^{\text{intra}} = \sum_{j=i}^{L_t} \lambda_{\text{intra}}^{j-i} \delta_{t,j}.\tag{18}$$

The final allocation solves a constrained least-squares projection:

$$\mathbf{A}_t^{\text{MAE}} = \underset{\mathbf{w} \in \mathbb{R}^{L_t}}{\operatorname{argmin}} \frac{1}{2}\|\mathbf{w} - \beta\mathbf{A}_t^{\text{intra}}\|_2^2, \quad \text{s.t.} \quad \frac{1}{L_t}\sum_{i=1}^{L_t} w_i = A_t^{\text{turn}},\tag{19}$$

yielding the closed-form solution:

$$\boxed{A_{t,i}^{\text{MAE}} = A_t^{\text{turn}} + \beta\left(A_{t,i}^{\text{intra}} - \overline{A}_t^{\text{intra}}\right).}\tag{20}$$

Turn-Level Importance Correction

The token and complete-turn importance ratios are:

$$\rho_{t,i}(\theta) = \frac{\pi_\theta(a_{t,i} \mid h_{t,i})}{\pi(a_{t,i} \mid h_{t,i})}, \qquad \rho_t(\theta) = \prod_{i=1}^{L_t} \rho_{t,i}(\theta).\tag{22}$$

To retain token-specific credit under complete-turn weighting, a stop-gradient construction is used:

$$\tilde{\rho}_{t,i}(\theta) = \rho_{t,i}(\theta)\operatorname{sg}\left[\prod_{j \neq i} \rho_{t,j}(\theta)\right],\tag{23}$$

with the gradient property:

$$\nabla_\theta \tilde{\rho}_{t,i}(\theta) = \rho_t(\theta) \nabla_\theta \log \pi_\theta(a_{t,i} \mid h_{t,i}).\tag{24}$$

The final MPO update uses a clipped surrogate:

$$\mathcal{L}_{\text{MPO}}(\theta) = \hat{\mathbb{E}}_{t,i}\left[\min\left(\tilde{\rho}_{t,i}(\theta)\hat{A}_{t,i}^{\text{MAE}}, \operatorname{clip}(\tilde{\rho}_{t,i}(\theta), 1-\epsilon, 1+\epsilon)\hat{A}_{t,i}^{\text{MAE}}\right)\right].\tag{25}$$

Empirical Validation / Results

Experimental Setup

• Models: Qwen3-4B-Thinking and Qwen3-30B-A3B-Instruct
• Tasks: Search (BrowseComp-Plus, GAIA, HLE), Claw-Eval, Automation Bench
• Baselines: GRPO, GSPO, PPO, VAPO, SAO, and MPO w/ token IS (ablation)

Main Results (Table 1)

Qwen3-4B-Thinking:

Qwen3-30B-A3B-Instruct:

Ablation Studies

Turn-scale propagation ($\lambda$): Optimal at $\lambda = 0.80$; increasing to 0.85 or decreasing to 0.60 both hurt performance.

Within-turn trace ($\lambda_{\text{intra}}$): Optimal at $\lambda_{\text{intra}} = 0.95$; both 0.90 and 1.00 reduce all scores.

Local refinement strength ($\beta$):

Test-Time Scaling

MPO maintains advantages at larger sampling budgets (pass@k for k ∈ {1, 2, 4, 8, 16, 32}). On 4B, MPO's 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 respectively.

Theoretical and Practical Implications

Theoretical Contributions

1. Turn-level improvement bound: Establishes that macro-action surrogates with complete-turn importance ratios provide valid monotonic improvement guarantees for LLM agents.

2. Token surrogate penalty: Identifies that token-level surrogates incur an additional within-turn distribution-shift penalty of $2(T+1)A_{\max}^{\text{tok}}L_{\max}(L_{\max}-1)\alpha_T^2$ under the same KL constraints.

3. Multiscale credit structure: Demonstrates empirically that critic values are primarily structured at turn boundaries (89.5% of variance between turns), motivating the turn-anchored approach.

Practical Implications

1. Distinct credit horizons: Interaction and generation require different decay rates—$\lambda$ decays per interaction while $\lambda_{\text{intra}}$ decays per token, with optimal values of 0.80 and 0.95 respectively.

2. Mean-preserving allocation: The centered local correction preserves turn-level credit while refining token distinctions, balancing long-range credit propagation with fine-grained allocation.

3. Stop-gradient importance weighting: Enables complete-turn probability correction while retaining token-specific gradients, aligning credit assignment with the actual interaction objective.

Conclusion

MPO organizes LLM-agent policy optimization across two scales: complete turns (macro actions) and tokens (generation steps). The MAE component combines cross-turn credit propagation with mean-preserving local allocation, while the update aligns complete-turn importance correction with token-specific gradients. Key findings include:

• Turn-level credit propagation with moderate local refinement ($\beta = 1.0$) provides the best performance
• Distinct decay rates for interactions ($\lambda = 0.80$) and tokens ($\lambda_{\text{intra}} = 0.95$) are essential
• Turn-level importance weighting outperforms token-level weighting, confirming the theoretical analysis
• Benefits persist under increased sampling (pass@k scaling)

Future directions suggested by this work include exploring adaptive multiscale coefficients, extending the framework to asynchronous training settings, and investigating whether the turn-token distinction generalizes to other hierarchical action structures in agentic systems.