Cost-Aware Policy Optimization (CAPO)

• CAPO is a framework that explicitly incorporates heterogeneous evaluation costs into policy optimization to achieve target performance with minimized total cost.
• It employs an importance sampling rule, pᵢ ∝ Gᵢ/√cᵢ, to adjust sample weights and reduce computational tokens while preserving or boosting accuracy.
• The approach extends to LLM policy optimization and opportunistic spectrum access, illustrating optimal cost-performance trade-offs in both reinforcement learning and control settings.

Cost-Aware Policy Optimization (CAPO) denotes policy-optimization methods in which cost is an explicit optimization variable rather than a secondary implementation concern. In the formulation instantiated by "Cost-Aware Learning" (Mohri et al., 30 Apr 2026), the objective is to reach a target error or benchmark performance while minimizing total cost, with cost arising from heterogeneous component evaluations in finite-sum optimization and, in LLM reinforcement learning, from sequence-length-dependent policy-gradient computation. A broader cost-aware control instantiation appears in opportunistic spectrum access, where the policy maximizes cumulative expected net reward, i.e., reward minus sensing and transmission costs, via structured threshold rules (Gan et al., 2018). The acronym is not unique; several papers use CAPO for unrelated methods, so the cost-aware usage requires terminological disambiguation.

1. Terminology and scope

In current arXiv usage, CAPO is an overloaded acronym. In the cost-aware sense relevant here, the term refers to optimizing policy updates under heterogeneous costs, especially when those costs are measurable and can be incorporated into sampling or decision rules. In "Cost-Aware Learning," the reinforcement-learning instantiation is called "Cost-Aware GRPO (CA-GRPO)," and the paper explicitly presents it as an instance of Cost-Aware Policy Optimization (Mohri et al., 30 Apr 2026).

Meaning of CAPO	Source	Scope
Cost-Aware Policy Optimization	(Mohri et al., 30 Apr 2026)	Cost-aware sampling and importance weighting in policy optimization
Credit Assignment Policy Optimization	(Xie et al., 4 Aug 2025)	Verifiable token-level credit assignment for reasoning LLMs
Consensus Aggregation for Policy Optimization	(Su et al., 13 Mar 2026)	Width-over-depth PPO consensus aggregation
Calibration-Aware Policy Optimization	(Wang et al., 14 Apr 2026)	AUC-consistent calibration optimization for reasoning LLMs

This terminological overlap matters because the cost-aware variant does not target token-level process rewards, PPO consensus aggregation, or relative calibration. Its central concern is the total cost required to achieve a target optimization error or policy-improvement threshold.

2. Finite-sum cost-aware optimization framework

The formal starting point is a finite-sum objective

$f(x) = \frac{1}{n} \sum_{i=1}^n f_i(x),$

where each component $f_i$ is convex and $G_i$-Lipschitz, and each component has an evaluation cost $c_i \ge 0$. Querying sample $i$ at point $x$ returns the exact gradient $\nabla f_i(x)$ and incurs cost $c_i$. The stated goal is to obtain $\hat x$ such that

$\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$

while minimizing total expected cost (Mohri et al., 30 Apr 2026).

The corresponding stochastic method is Cost-Aware SGD, an importance-weighted projected SGD scheme. With sampling distribution $f_i$0, the stochastic gradient estimator is

$f_i$1

which is unbiased:

$f_i$2

Its second moment is bounded by

$f_i$3

and the expected per-iteration cost is

$f_i$4

If $f_i$5 is the number of iterations required to reach error $f_i$6, then the total expected cost is

$f_i$7

For general convex functions, the paper states

$f_i$8

so achieving error $f_i$9 requires

$G_i$0

and hence

$G_i$1

For $G_i$2-strongly convex functions with $G_i$3, the paper gives

$G_i$4

In both regimes, minimizing total cost reduces to minimizing the product

$G_i$5

Theorem 4.1 identifies the optimal sampling distribution:

$G_i$6

Under this distribution, the minimum expected cost becomes

$G_i$7

for the convex case, and

$G_i$8

for the strongly convex case. The comparison point is traditional variance-only importance sampling, which uses $G_i$9 and ignores costs. Theorem 4.2 states that the optimal cost-aware strategy is never worse than uniform or variance-only sampling in total cost to reach $c_i \ge 0$0-accuracy.

A plausible implication is that CAPO changes the optimization target from “minimize iteration count” to “minimize cost-adjusted convergence time.” In the LLM setting, this alters which trajectories are most valuable to revisit during policy updates.

3. Optimality, proxy quality, and subset selection

The theory in "Cost-Aware Learning" does not stop at the optimal distribution. It also characterizes lower bounds, proxy error, and biased subset restriction (Mohri et al., 30 Apr 2026).

The lower-bound result states, informally, that for any algorithm achieving $c_i \ge 0$1, there exists a convex, $c_i \ge 0$2-Lipschitz finite-sum instance over $c_i \ge 0$3 such that the expected total query cost is at least

$c_i \ge 0$4

for any subset $c_i \ge 0$5 satisfying the stated cost-uniformity condition. The paper’s interpretation is that the dependence on $c_i \ge 0$6 in the upper bounds is fundamental in heterogeneous-cost regimes.

Because exact $c_i \ge 0$7 values are often unavailable, the paper gives an exact proxy sub-optimality identity. Let $c_i \ge 0$8 be the cost-biased version of a sampling distribution. Then Theorem 4.3 states

$c_i \ge 0$9

Thus the gap between a practical sampler and the optimal sampler is determined exactly by the Pearson $i$0-divergence between their cost-biased distributions. Theorem 4.4 further analyzes noisy proxies $i$1 and gives an approximate dependence on the proxy–truth Pearson correlation $i$2.

The subset-selection analysis introduces a bias–cost trade-off. Restricting optimization to a subset $i$3 yields

$i$4

which generally converges to $i$5. With bias upper bound $i$6, the convex-case guarantee becomes

$i$7

Theorem 5.1 states that, for a fixed bias budget $i$8 with requirement $i$9, the optimal subset is equivalent to a Min-Knapsack problem:

$x$0

with item cost $x$1, item value $x$2, and

$x$3

The greedy approximation sorts by density $x$4 and therefore selects the cheapest samples first.

This suggests a two-level CAPO picture: first optimize sampling probabilities inside the full dataset, and then, under stronger budget pressure, restrict the support itself through an explicit bias-aware selection rule.

4. CAPO in GRPO-based LLM policy optimization

The LLM instantiation applies the cost-aware finite-sum analysis to Group Relative Policy Optimization. GRPO alternates between data collection and policy-gradient updates, uses a token-level PPO-like objective with clipping and KL regularization, and normalizes advantages as

$x$5

The empirical finite-sum GRPO objective is written as

$x$6

with rollout-level contributions $x$7 defined by the clipped ratio term and the KL regularizer (Mohri et al., 30 Apr 2026).

The cost model is sequence-length based:

$x$8

namely the number of tokens processed during the policy-gradient step, which the paper states is proportional to FLOPs. This converts policy optimization into a heterogeneous-cost finite-sum problem, because longer trajectories are more expensive to reuse during the update stage.

Cost-Aware GRPO instantiates CAPO by importance sampling rollouts during policy updates using the cost-aware rule. Since exact per-sample gradient norms are not computed, the paper uses the magnitude of the normalized advantage, $x$9, as a proxy for $\nabla f_i(x)$0. Sampling probabilities for rollout $\nabla f_i(x)$1 in the step dataset $\nabla f_i(x)$2 are therefore set to

$\nabla f_i(x)$3

with an optional smoothed distribution

$\nabla f_i(x)$4

where $\nabla f_i(x)$5 is uniform and $\nabla f_i(x)$6. The bias induced by nonuniform sampling is corrected by the importance-weighted objective

$\nabla f_i(x)$7

An implementation detail recentered the importance weights around $\nabla f_i(x)$8 in order to maintain stable gradient norms.

Algorithm 2 proceeds in three phases per iteration: rollout collection, computation of the sampling distribution from $\nabla f_i(x)$9, and multinomial mini-batch sampling for policy updates. The practical guidance is explicit that data collection is not reweighted; only the policy-gradient stage is cost-aware. The same section also notes that smoothing values such as $c_i$0, $c_i$1, and $c_i$2 were explored, and that the method is robust across GRPO, GRPO+ZVF, and CISPO.

5. Empirical results in reasoning LLMs

The empirical study uses Qwen2.5-Math-1.5B-Instruct and Qwen3-8B Base, trains on the DAPO dataset with verifiable math rewards, and evaluates on MATH500, AMC, GSM8K, and AIME1983–2024 using pass@1 or mean@32. The implementation uses Verl, vLLM for inference, and FlashAttention for efficiency. For Qwen3-8B runs, the hardware is 8× H100 GPUs with approximately 5 days per run; the 1.5B runs take approximately 1–1.5 days (Mohri et al., 30 Apr 2026).

The main reported outcome is a reduction in cumulative tokens used in policy optimization while matching or exceeding baseline accuracy. The paper also reports that the proxy $c_i$3 is empirically well aligned with true gradient norms: Pearson correlations are approximately $c_i$4–$c_i$5, and cost-biased $c_i$6-divergences are approximately $c_i$7–$c_i$8 across training steps.

Setting	Token savings	Accuracy summary
Qwen2.5-Math-1.5B-Instruct, standard GRPO on AIME	$c_i$9 and $\hat x$0 match baseline peak using 28% and 32% fewer tokens	Both ultimately exceed baseline by approximately 5% peak accuracy
Qwen2.5-Math-1.5B-Instruct, GRPO+ZVF on AIME	$\hat x$1 achieves baseline peak with 13% fewer tokens	Exceeds baseline by approximately 5%
Qwen3-8B, AIME	$\hat x$2 reaches baseline peak using 51% fewer tokens; at 99% of peak accuracy, 17% fewer tokens	Peak accuracy preserved or slightly improved
Qwen3-8B, AMC	$\hat x$3 reaches baseline maximum using 47% fewer tokens	Matches baseline maximum

The reported peak accuracy averages further emphasize that cost reduction does not require an accuracy sacrifice. For the 1.5B model under GRPO, the best averaged checkpoint over four benchmarks is 71.2% with no sampling, 74.0% with $\hat x$4, and 72.9% with $\hat x$5. Under GRPO+ZVF, the same comparison is 70.6%, 69.8%, and 72.5%. For the 8B model under GRPO, the corresponding values are 78.5% with no sampling, 78.8% with $\hat x$6, and 78.5% with both $\hat x$7 and $\hat x$8.

The ablations are important for interpretation. Under CISPO, $\hat x$9 and $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$0 require 31% and 34% fewer tokens to reach baseline peak, respectively. By contrast, the length-only distribution $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$1–LEN, i.e., $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$2, underperforms. This directly supports the claim that CAPO is not merely “prefer shorter sequences”; it is signal-aware sampling proportional to estimated information-to-cost ratio.

6. Broader cost-aware control formulation

A distinct but conceptually related instantiation appears in cost-aware opportunistic spectrum access, where the objective is to maximize cumulative expected net reward over $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$3 frames under random sensing costs, random transmission costs, and stochastic transmission rewards (Gan et al., 2018). Each frame has a sensing phase followed by a transmission phase. Channel $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$4 is idle with probability $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$5, and if transmission occurs on an idle channel the reward is a random variable with mean $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$6; sensing and transmission incur random costs with means $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$7 and $\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$8.

The per-frame decision process admits a Bellman equation

$\mathbb{E}[f(\hat x) - f(x^*)] \le \varepsilon$9

When the last sensed channel is busy, the value recursion becomes

$f_i$00

with $f_i$01. The resulting optimal offline policy has a recursive double-threshold structure. For each position $f_i$02, there are thresholds

$f_i$03

such that:

• if $f_i$04, the policy guesses and transmits without sensing;
• if $f_i$05, it senses channel $f_i$06;
• if $f_i$07, it quits the frame.

The online version estimates $f_i$08, $f_i$09, $f_i$10, and $f_i$11, uses an exploration threshold $f_i$12, and alternates between exploration and plug-in exploitation. The regret guarantee is

$f_i$13

with the stated constants

$f_i$14

and the paper also establishes a matching $f_i$15 lower bound under mild conditions.

This earlier line of work differs from LLM CAPO in mechanism, but not in principle: in both cases, costs affect the optimal policy directly rather than appearing only as after-the-fact efficiency metrics.

7. Limitations and recurring misconceptions

The LLM-oriented CAPO framework comes with several explicit caveats. The underlying theory assumes convex or strongly convex finite-sum objectives, whereas LLM policy optimization is non-convex and non-stationary because rollout pools change across training rounds. The paper argues that the $f_i$16 objective remains structurally optimal due to variance-controlled iteration complexity, but this is not a convex-style global guarantee (Mohri et al., 30 Apr 2026).

The method also relies on gradient proxies rather than exact gradient norms. In the reported math-reasoning setting, $f_i$17 works well, but the paper notes that long-horizon credit assignment and sparse rewards can make advantage magnitudes noisy; smoothing and ZVF help, but more robust proxies may be needed. Distribution shift across training rounds further means that guarantees are per rollout pool rather than global. Subset selection introduces bias and therefore requires a carefully chosen bias budget. A related misconception is to treat length-only sampling as sufficient; the reported underperformance of $f_i$18–LEN shows that explicit signal terms remain necessary.

Finally, CAPO should not be treated as a unique algorithm name. In recent work, the acronym also denotes Credit Assignment Policy Optimization, Consensus Aggregation for Policy Optimization, and Calibration-Aware Policy Optimization (Xie et al., 4 Aug 2025, Su et al., 13 Mar 2026, Wang et al., 14 Apr 2026). In the cost-aware sense, the defining feature is not verifiable process supervision, optimizer-path consensus, or AUC-consistent calibration; it is the direct optimization of performance relative to heterogeneous cost. A plausible implication is that future CAPO variants for RLHF will be judged less by raw iteration counts than by how effectively they convert compute, measured in tokens or FLOPs, into policy improvement.