Step-Aware Policy Optimization (SAPO)

• Step-Aware Policy Optimization is a reinforcement learning framework that focuses on intermediate decision steps for more precise credit assignment.
• It reformulates policy improvement by leveraging step-level optimization techniques like TD residuals and GRPO-style updates to handle multi-stage computations.
• Empirical findings reveal enhanced efficiency and performance in diverse applications, including diffusion language models and agentic search tasks.

Step-aware policy optimization denotes a set of reinforcement-learning formulations that align policy improvement with an intermediate decision unit rather than treating only full trajectories or individual tokens as the primary locus of credit assignment. In current literature, the label is used in several closely related ways: as the exact title “Step-Aware Policy Optimization” for reasoning in diffusion LLMs (Xie et al., 2 Oct 2025), as “Step-Aligned Policy Optimization” for agentic reinforcement learning (Wang et al., 20 Apr 2026), and as “Step Advantage Policy Optimization” for graph-grounded agentic search (Liu et al., 28 May 2026). Across these formulations, the common objective is to reduce the mismatch between sparse outcome rewards and the structured, multi-stage computations executed by modern language-model agents.

1. Terminology and acronym ambiguity

The phrase has substantial naming ambiguity. The exact name “Step-Aware Policy Optimization” appears in work on diffusion LLMs, where the denoising trajectory is treated as a latent reasoning process (Xie et al., 2 Oct 2025). Closely related agentic-RL work instead uses the name StepPO, “Step-Aligned Policy Optimization,” and argues that the proper action abstraction for LLM agents is the complete interaction step rather than the token (Wang et al., 20 Apr 2026). A further adjacent formulation, “Step Advantage Policy Optimization,” couples step-level advantages to graph-based process rewards in agentic search (Liu et al., 28 May 2026).

The acronym SAPO is also overloaded by unrelated methods. In the same period, SAPO denotes “Soft Adaptive Policy Optimization” for smooth-gated off-policy LLM alignment (Gao et al., 25 Nov 2025), “Search Agent Policy Optimization” for conditional token-level KL stabilization in tool-based search RL (Li et al., 10 Mar 2026), “Self-Augmented Preference Optimization” for off-policy preference learning (Yin et al., 2024), and “Soft Analytic Policy Optimization” for differentiable multiphysics simulation (Xing et al., 2024). Consequently, the term must be interpreted from paper context rather than acronym alone.

2. Step-level formulation in agentic reinforcement learning

The clearest formalization of step-aware optimization in agentic RL is StepPO. Its starting point is the claim that the token-centric Markov decision process inherited from autoregressive language modeling is inadequate for multi-turn agent training, where the model alternates between reasoning, tool use, environment transitions, and context updates (Wang et al., 20 Apr 2026). StepPO therefore advances a step-level MDP in which the trajectory is

$$\tau = \{(s_t, a_t, r_t, s_{t+1})\}_{t=0}^{T-1}, \qquad J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}\left[\sum_{t=0}^{T-1}\gamma^t r_t\right].$$

Here $a_t$ is not a single token but a complete interaction action: a token sequence, a structured tool call, or a mixed reasoning-and-action response.

This formulation does not discard token factorization. A step action may still be written as

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$

with variable step length $L_t$. The policy ratio is then elevated from token to step granularity: $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$ The optimization counterpart is step-level credit assignment via TD residuals and GAE: $$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t), \qquad A_t = \sum_{l=0}^{T-t-1} (\gamma \lambda)^l \delta_{t+l}.$$ The actor update is a PPO-style clipped surrogate at step granularity,

$$\mathcal{L}_{\mathrm{actor}}(\theta)= \mathbb{E}\left[ \min\left( w_t(\theta)A_t,\, \operatorname{clip}(w_t(\theta),1-\epsilon_{\mathrm{low}},1+\epsilon_{\mathrm{high}})A_t \right) \right].$$

The significance of this construction is conceptual as much as algorithmic. StepPO positions the interaction step as the semantic decision unit, and aligns MDP formulation, credit assignment, replay structure, and systems design to that unit. A plausible implication is that it imports a semi-MDP-like view into LLM agents, although the paper does not formally adopt semi-MDP notation (Wang et al., 20 Apr 2026).

3. Step-level rewards and advantages in agentic search

In agentic search, step-aware optimization has been developed through explicit process rewards. “Beyond Trajectory Rewards: Step-level Credit Assignment for Agentic Search via Graph Modeling” introduces Graph-Distance Contribution Reward (GDCR) and Step Advantage Policy Optimization (SAPO) (Liu et al., 28 May 2026). The paper views world knowledge as a latent world graph $\mathcal{G}^{\mathrm{world}}$ and each query-specific task as search in a latent task graph $\mathcal{G}_q$, with answer node $v_q^\ast$. A training-time proxy graph $a_t$0 supports a distance-based contribution score

$a_t$1

GDCR rewards both newly cited entities and newly retrieved entities. With cumulative cited set $a_t$2 and observed set $a_t$3,

$a_t$4

the step reward is

$a_t$5

SAPO then converts these process rewards into bounded step-level advantages,

$a_t$6

and combines them with trajectory-level outcome advantage: $a_t$7 The resulting policy objective is a GRPO-style clipped surrogate in which each token inherits the advantage of its containing step.

This formulation is explicitly intended to overcome the inadequacy of trajectory-level outcome rewards for long-horizon information seeking. It also preserves efficiency: for Qwen3-8B, the paper reports average training cost per step of 13.13 min for GRPO and 13.26 min for SAPO, compared with 16.25 min for ARPO and an estimated $a_t$8 min for Tree-GRPO with full tree expansion and $a_t$9 (Liu et al., 28 May 2026).

4. Selective and success-rate-aware step optimization

A second line of work makes step-aware optimization selective rather than uniform. STAPO, “Selective Trajectory-Aware Policy Optimization,” addresses trajectory neglect in long-horizon LLM agents by identifying anomalous intermediate steps through normalized entropy and applying additional step-specific shaping only there (Qi et al., 6 Jul 2026). For action $a_t = (y_{t,1}, \dots, y_{t,L_t}),$0, step entropy is

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$1

and state-normalized entropy is

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$2

An IQR test yields an outlier mask

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$3

Only outlier steps receive the selective term

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$4

where $a_t = (y_{t,1}, \dots, y_{t,L_t}),$5 is a trajectory-aware reward based on KL divergence between full and trajectory-blind prompts, and $a_t = (y_{t,1}, \dots, y_{t,L_t}),$6 is a trajectory-independent penalty against a reference policy. STAPO thereby combines step selection, step-level uncertainty calibration, and per-step auxiliary optimization.

STEP, “Success-rate-aware Trajectory-Efficient Policy Optimization,” shifts emphasis from selective step shaping to task-aware sampling and step-level reuse (Chen et al., 17 Nov 2025). It maintains a smoothed per-task success-rate record $a_t = (y_{t,1}, \dots, y_{t,L_t}),$7, applies replacement probability

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$8

and updates success statistics via

$a_t = (y_{t,1}, \dots, y_{t,L_t}),$9

For successful trajectories only, it assigns

$L_t$0

then augments low-success tasks with local step-level GRPO groups

$L_t$1

and final step advantage

$L_t$2

Both STAPO and STEP are not titled SAPO, but each instantiates step-aware policy optimization in the substantive sense of localizing credit and optimization pressure to selected intermediate decisions.

5. Step-aware optimization for diffusion LLMs

The exact title “Step-Aware Policy Optimization” appears in diffusion-language-model reasoning (Xie et al., 2 Oct 2025). The target setting is mask-based diffusion LLMs, where generation proceeds by iterative denoising rather than autoregressive token emission. The paper argues that outcome-only RL induces unstructured refinement, a failure mode in which denoising steps do not contribute meaningfully to the solution. Its theoretical framework models reasoning as a hierarchical selection process,

$L_t$3

with observable response $L_t$4 generated by inverting a latent hierarchy of simpler subproblems. Under smoothness, sparsity, and structural-diversity assumptions, the paper states that the latent reasoning hierarchy is identifiable up to level-wise permutations and invertible transformations (Xie et al., 2 Oct 2025).

SAPO operationalizes this theory with a process reward defined on denoising intervals. For two timesteps $L_t$5, and rollout completions from intermediate states $L_t$6 and $L_t$7, the process reward is

$L_t$8

In practice the efficient special case $L_t$9 is used. This quantity measures whether the denoising interval from $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$0 to $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$1 increases eventual success probability. The total advantage is then

$$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$2

where $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$3 is the baseline GRPO-style advantage.

The conceptual importance of this construction is that it treats denoising steps as proxies for latent reasoning stages. The paper does not claim an exact one-to-one correspondence between denoising steps and logical steps; instead, the denoising interval is a proxy for resolving a set of logical constraints. This yields a form of step awareness specific to non-autoregressive generation.

6. Empirical support, systems consequences, and open issues

Empirical support is uneven but substantial. StepPO reports only preliminary evidence on HotpotQA with Qwen2.5-3B-Instruct, under a controlled comparison in which it consistently outperforms token-level PPO throughout training, remains above the token-level baseline across most of the run, reaches a noticeably higher peak, and maintains a stronger plateau in middle and later training (Wang et al., 20 Apr 2026). The paper also specifies a stepwise rollout scheme, prompt budget 10,240, response budget 1,024, and $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$4, $$w_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)} =\prod_{i=1}^{L_t}\frac{\pi_\theta(y_{t,i}\mid s_t,y_{t,<i})}{\pi_{\theta_{\mathrm{old}}}(y_{t,i}\mid s_t,y_{t,<i})}.$$5.

Graph-based SAPO reports stronger quantitative gains. On Qwen3-30B-A3B-thinking, the transition from GRPO to SAPO changes benchmark scores from 33.2 / 14.9 / 67.0 / 62.1 to 45.7 / 42.8 / 75.0 / 70.9 across BrowseComp-ZH / BrowseComp / xbench-DS / GAIA (Liu et al., 28 May 2026). STAPO reports improvements over GiGPO on long-horizon agent tasks: for Qwen2.5-7B-Instruct, success on ALFWorld rises from 90.8% to 96.9%, and success on WebShop rises from 72.8% to 77.6% (Qi et al., 6 Jul 2026). STEP reports OSWorld overall improvement from 18.9 with T-GRPO and 21.1 with GiGRPO to 23.8 with STEP, and AndroidWorld overall improvement from 31.0 and 34.0 to 45.7, while reducing average training time per step from 45.67 min/step for T-GRPO to 26.25 min/step (Chen et al., 17 Nov 2025). Diffusion SAPO shows especially large gains on COUNTDOWN, where scores change from 33.2 / 31.3 / 37.1 for diffu-GRPO to 51.6 / 52.0 / 56.3 across generation lengths 128 / 256 / 512 (Xie et al., 2 Oct 2025).

Several limitations recur. StepPO leaves critic loss, KL regularization, and many optimization details unspecified, and its empirical support is preliminary (Wang et al., 20 Apr 2026). Graph-based SAPO depends on the quality of the training-time ER graph and answer-node annotation, and does not directly apply to open-ended tasks without a clear endpoint (Liu et al., 28 May 2026). STAPO assumes sufficient within-batch support for each anchor state; if a state is visited only once, normalized entropy degenerates toward raw entropy (Qi et al., 6 Jul 2026). STEP discards failed trajectories entirely, which avoids penalizing correct intermediate actions in failed runs but also discards potentially useful sub-trajectories (Chen et al., 17 Nov 2025). Diffusion SAPO explicitly relies on the mean-field assumption in diffu-GRPO for estimating sequence likelihoods, and its gains on MATH are less uniform than on COUNTDOWN or SUDOKU (Xie et al., 2 Oct 2025).

A final point of interpretation is terminological rather than algorithmic. Because “SAPO” simultaneously denotes step-aware, step-advantage, soft-adaptive, search-agent, self-augmented, and soft-analytic formulations in different papers, the literature is best read through the underlying optimization unit. In the step-aware sense, the defining feature is consistent: policy optimization is reorganized around semantically meaningful intermediate decisions—interaction steps, search steps, denoising intervals, or selected outlier steps—so that reward propagation and trust-region control better match the actual structure of reasoning and agent behavior.