---
title: Step-Aware Policy Optimization (SAPO)
url: https://www.emergentmind.com/topics/step-aware-policy-optimization-sapo
type: topic
---

# Step-Aware Policy Optimization (SAPO)

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 language models [2510.01544], as “Step-Aligned Policy Optimization” for agentic reinforcement learning [2604.18401], and as “Step Advantage Policy Optimization” for graph-grounded agentic search [2605.29697]. 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 large language models, where the denoising trajectory is treated as a latent reasoning process [2510.01544]. 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 [2604.18401]. A further adjacent formulation, “Step Advantage Policy Optimization,” couples step-level advantages to graph-based process rewards in agentic search [2605.29697].

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 [2511.20347], “Search Agent Policy Optimization” for conditional token-level KL stabilization in tool-based search RL [2603.10069], “Self-Augmented Preference Optimization” for off-policy preference learning [2405.20830], and “Soft Analytic Policy Optimization” for differentiable multiphysics simulation [2412.12089]. 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 [2604.18401]. 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 [2604.18401].

## 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) [2605.29697]. 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 \(\hat{\mathcal{G}}_q\) supports a distance-based contribution score
\[
d_q(v, v_q^\ast) = \min_{p \in \hat{\mathcal{P}}(v, v_q^\ast)} |p|,
\qquad
c_q(v) = k^{-d_q(v, v_q^\ast)}.
\]

GDCR rewards both newly cited entities and newly retrieved entities. With cumulative cited set \(T_t\) and observed set \(O_t\),
\[
\Delta C_t = T_t \setminus T_{t-1},
\qquad
\Delta O_t = O_t \setminus O_{t-1},
\]
the step reward is
\[
r_t^{\mathrm{cite}} = \sum_{v \in \Delta C_t} c_q(v),
\qquad
r_t^{\mathrm{ret}} = \sum_{v \in \Delta O_t} c_q(v),
\qquad
r_{t,g} = r_t^{\mathrm{cite}} + r_t^{\mathrm{ret}}.
\]
SAPO then converts these process rewards into bounded step-level advantages,
\[
A^{(i)}_{t,g} =
\mathrm{clip}\left(
\frac{r^{(i)}_{t,g}-\mathrm{mean}_{t=1}^{T_i}(r^{(i)}_{t,g})}
{\mathrm{std}_{t=1}^{T_i}(r^{(i)}_{t,g})+\varepsilon},
-1,1
\right),
\]
and combines them with trajectory-level outcome advantage:
\[
A^{(i)}_{t} = \hat{A}^{(i)}_o + \lambda |\hat{A}^{(i)}_o| A^{(i)}_{t,g}.
\]
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 **\(\sim 2107\) min** for Tree-GRPO with full tree expansion and \(B=2\) [2605.29697].

## 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 [2607.04963]. For action \(a_t^{(i)}\), step entropy is
\[
H(a_t^{(i)}) = - \frac{1}{M} \sum_{k=1}^{M} \sum_{v \in \mathcal{V}}
p(v \mid c_t^{(i)}, w_{<k}) \log p(v \mid c_t^{(i)}, w_{<k}),
\]
and state-normalized entropy is
\[
H_n(a_t^{(i)}) =
\frac{
H(a_t^{(i)}) -
\mu(\{H(a_t^{(j)})\})
}{
\sigma(\{H(a_t^{(j)})\}) + \epsilon
}.
\]
An IQR test yields an outlier mask
\[
M_t(a_t^{(i)}) = \mathbb{I}\!\left(H_n(a_t^{(i)}) \notin \mathcal{R}\right).
\]
Only outlier steps receive the selective term
\[
\mathcal{J}_{\mathrm{sel}}^{(i,t)}
=
M_t \cdot \bigl( \alpha R_{TA}(c_t^{(i)}, a_t^{(i)}) - \gamma P_{TI}(c_t^{(i)}, a_t^{(i)}) \bigr),
\]
where \(R_{TA}\) is a trajectory-aware reward based on KL divergence between full and trajectory-blind prompts, and \(P_{TI}\) 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 [2511.13091]. It maintains a smoothed per-task success-rate record \(s_i\), applies replacement probability
\[
P_{\mathrm{rep}}(s_i) = \frac{1}{1 + \exp(-K(s_i - s_0))},
\]
and updates success statistics via
\[
\tilde{N}_i^T = N_i^T + \alpha \tilde{N}_i^{T-1},
\qquad
\tilde{U}_i^T = U_i^T + \alpha \tilde{U}_i^{T-1},
\qquad
s_i = \frac{\tilde{U}_i^T}{\tilde{N}_i^T}.
\]
For successful trajectories only, it assigns
\[
\mathrm{Adv}(T_{i,j}) = (1 - s_i)\cdot R_{T_{i,j}},
\qquad
\mathrm{Adv}(S_{T,t}) = \mathrm{Adv}(T),
\]
then augments low-success tasks with local step-level GRPO groups
\[
G_S = \{S_k = (S, A_k) \mid k=0, \ldots, n,\; A_0 = A\},
\]
and final step advantage
\[
\mathrm{Adv}_{\mathrm{final}}(S_k) =
\mathrm{Adv}(S)\cdot \mathrm{Adv}_{\mathrm{aug}}(S_k).
\]
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 language models

The exact title “Step-Aware Policy Optimization” appears in diffusion-language-model reasoning [2510.01544]. The target setting is mask-based diffusion language models, 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,
\[
\mS_l := g_{\mS_l}(\mS_{l+1}),
\quad l = 0,\dots,L-1,
\quad \mS_0 := \mQ,
\]
with observable response \(\mR\) 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 [2510.01544].

SAPO operationalizes this theory with a process reward defined on denoising intervals. For two timesteps \(0 \le t_1 < t_2 \le T\), and rollout completions from intermediate states \(x_{t_1}\) and \(x_{t_2}\), the process reward is
\[
R_{\text{process}}(t_1, t_2) =
\frac{1}{N_1}\sum_{j=1}^{N_1} \mathbf{1}[\mR^{(j)}(x_{t_1})]
-
\frac{1}{N_2}\sum_{j=1}^{N_2} \mathbf{1}[\mR^{(j)}(x_{t_2})].
\]
In practice the efficient special case \(t_2=T\) is used. This quantity measures whether the denoising interval from \(t_2\) to \(t_1\) increases eventual success probability. The total advantage is then
\[
A^{\text{total}}_i = A_i + \mathbf{1}[A_i > 0] \cdot R_{\text{process}},
\]
where \(A_i\) 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 [2604.18401]. The paper also specifies a stepwise rollout scheme, prompt budget **10,240**, response budget **1,024**, and \(\gamma=0.99\), \(\lambda=1.0\).

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** [2605.29697]. 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%** [2607.04963]. 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** [2511.13091]. 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** [2510.01544].

Several limitations recur. StepPO leaves critic loss, KL regularization, and many optimization details unspecified, and its empirical support is preliminary [2604.18401]. 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 [2605.29697]. STAPO assumes sufficient within-batch support for each anchor state; if a state is visited only once, normalized entropy degenerates toward raw entropy [2607.04963]. STEP discards failed trajectories entirely, which avoids penalizing correct intermediate actions in failed runs but also discards potentially useful sub-trajectories [2511.13091]. 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** [2510.01544].

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.

Source: https://www.emergentmind.com/topics/step-aware-policy-optimization-sapo