💬 NLP

StepPO: Step-Aligned Policy Optimization for Agentic Reinforcement Learning

This paper introduces StepPO, a framework that advances Agentic Reinforcement Learning by shifting the optimization paradigm from token-level to step-level Markov Decision Processes, thereby aligning policy updates and credit assignment with the natural granularity of agent decisions to better handle multi-turn interactions, sparse rewards, and complex tool use.

Original authors: Daoyu Wang, Qingchuan Li, Mingyue Cheng, Jie Ouyang, Shuo Yu, Qi Liu, Enhong Chen

Published 2026-04-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Daoyu Wang, Qingchuan Li, Mingyue Cheng, Jie Ouyang, Shuo Yu, Qi Liu, Enhong Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Teaching a Robot to Play a Game

Imagine you are trying to teach a robot (a Large Language Model) how to play a complex video game, like The Legend of Zelda. The robot needs to explore a dungeon, find keys, defeat monsters, and solve puzzles to win.

In the past, when we trained AI, we mostly taught them to write sentences or answer questions in one go. We looked at their output letter by letter (or "token by token").

The Problem:
When the robot is playing a game, looking at it one letter at a time is like trying to teach someone how to drive a car by critiquing every single movement of their finger on the steering wheel.

"Your finger moved left too much."
"Your finger moved right too little."
"You pressed the gas pedal too hard."

This is Token-Level Learning. It's too noisy and granular. The robot doesn't know why it failed. Did it fail because it turned the wheel slightly wrong, or because it decided to drive into a wall in the first place?

The Solution: StepPO
The authors of this paper propose a new way to train these "Agentic" AI systems. They call it StepPO (Step-Aligned Policy Optimization).

Instead of watching the robot's finger movements, StepPO watches the robot's decisions.

Step 1: The robot sees a locked door.
Step 2: The robot decides to "Search for a key."
Step 3: The robot finds the key and opens the door.
Step 4: The robot gets a reward (points).

StepPO says: "Don't blame the robot for the specific letters it typed to say 'Search for a key.' Blame the decision to search for a key."

The Three Key Changes (The "How-To")

The paper argues that to make this work, we need to change three things in how we train AI, using a simple analogy of building a house.

1. The Blueprint (MDP Formulation)

Old Way (Token-Level): The blueprint is a list of every single brick laid. "Brick 1, Brick 2, Brick 3..." If you mess up Brick 42, the whole plan is confusing.
New Way (Step-Level): The blueprint is a list of tasks. "Lay the foundation," "Build the walls," "Put on the roof."
The Analogy: In StepPO, the "Action" isn't typing a word; it's completing a whole interaction round. Did the robot successfully call a tool? Did it successfully ask a question? That whole chunk is one "Step."

2. The Scorecard (Credit Assignment)

Old Way: If the house collapses at the end, the old system tries to figure out which single brick caused the collapse. It spreads the blame (or credit) across thousands of tiny bricks. This is messy and often wrong.
New Way: If the house collapses, we look at the tasks. "The foundation was weak," or "The roof was built too early."
The Analogy: StepPO gives the "reward" (the points) to the Step that made the difference. If the robot found the key in Step 2, and that led to winning the game in Step 10, StepPO says, "Great job on Step 2!" It connects the reward back to the specific decision that mattered, rather than the random words typed during that decision.

3. The Construction Site (Training Systems)

Old Way: The construction crew works in a chaotic line. They write everything down on a long scroll of paper. If they need to fix a mistake, they have to rewrite the whole scroll.
New Way: The crew uses digital blocks. Each "Step" is a self-contained block (State + Action + Reward). If a block is bad, you just swap that block out without ruining the rest of the house.
The Analogy: The paper mentions that real-world AI training needs a better "database" to store these blocks. Instead of a messy text file, they use a structured system where every "Step" is saved perfectly, so the AI can learn from its mistakes without getting confused by messy data.

Why Does This Matter? (The Results)

The authors tested this on a difficult question-answering game called HotpotQA (where the AI has to search through many documents to find an answer).

The Result: The AI trained with StepPO learned much faster and got better scores than the AI trained with the old "letter-by-letter" method.
Why? Because the AI learned to make better decisions (like "I should search for this specific fact") rather than just trying to guess the next letter correctly.

Summary in One Sentence

StepPO is a new training method that stops teaching AI how to type words and starts teaching it how to make decisions, by treating every interaction as a single, complete "step" rather than a long string of letters.

The "Takeaway" Metaphor

Imagine you are teaching a child to ride a bike.

Token-Level Training: You stand next to them and shout, "Pedal harder! Lean left! Turn the handlebars right! Pedal harder!" It's overwhelming and confusing.
StepPO Training: You wait until they finish a lap. Then you say, "Great job turning at the corner! Next time, try pedaling a bit faster before the hill." You focus on the milestones (the steps), not the muscle movements (the tokens).

This paper argues that for AI to become a true "Agent" that can do complex tasks, we must switch to the StepPO way of teaching.

1. Problem Statement

The paper addresses a fundamental granularity mismatch in current Agentic Reinforcement Learning (RL) systems.

Context: Large Language Model (LLM) agents are evolving from single-turn responders to multi-turn, long-horizon systems capable of tool use, planning, and environment interaction (e.g., OpenClaw, Claude Code).
The Mismatch: Traditional RL for LLMs (e.g., RLHF, PPO) relies on a token-level Markov Decision Process (MDP) where the "action" is generating a single token. However, agentic behavior is inherently step-level: an agent makes a decision (e.g., "call a tool," "search the web"), executes it, and receives an observation.
Consequences:
- Credit Assignment Noise: Propagating rewards at the token level is too fine-grained and noisy for long-horizon tasks, while trajectory-level assignment is too coarse to distinguish between specific intermediate decisions.
- Retokenization Drift: Storing agent interactions as text messages and re-tokenizing them for training introduces inconsistencies between the rollout (inference) and the training data, breaking the mathematical equivalence required for stable policy optimization.
- System Inefficiency: Current systems struggle to manage heterogeneous agents, long contexts, and asynchronous data collection efficiently because they lack a native "step" abstraction.

2. Methodology: StepPO

The authors propose StepPO, a framework that aligns the entire RL pipeline—from modeling to systems design—with the interaction step rather than the token.

A. Theoretical Shift: Step-Level MDP

Formulation: The paper redefines the MDP transition unit.
- Token-Level: State $s_i$ , Action $a_i$ (one token), Transition to $s_{i+1}$ .
- Step-Level: State $s_t$ (observation/context), Action $a_t$ (complete interaction round, e.g., a tool call + reasoning), Reward $r_t$ , Transition to $s_{t+1}$ (new environment state).
Action Definition: An action $a_t$ is a complete token sequence generated by the policy to execute a specific sub-task, not just a single token.

B. Trajectory Representation

To support the Step-Level MDP, the paper advocates for a Structured Step-Level Representation over text-space or flat token-space:

Structure: Trajectories are stored as a sequence of discrete tuples $(s_t, a_t, r_t)$ .
Consistency: This avoids "retokenization drift" by storing actions as raw token IDs (or structured objects) rather than decoded text that must be re-tokenized.
Flexibility: Allows for surgical context management (truncating/summarizing specific steps) without breaking the token-level consistency of other steps.

C. Optimization: Step-Level Credit Assignment

StepPO introduces a new objective function that aligns reward propagation with the step unit:

Step-Level Advantage: Instead of token-level GAE or trajectory-level advantage, the algorithm computes the Generalized Advantage Estimation (GAE) at the step level:
$\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$
$A_t = \sum_{l=0}^{T-t-1} (\gamma \lambda)^l \delta_{t+l}$
Policy Objective: The PPO surrogate loss is calculated based on the ratio of probabilities of the entire step (the full token sequence of the action) rather than individual tokens:
$w_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\theta_{old}}(a_t | s_t)}$
Benefit: This ensures that delayed rewards are attributed to the specific decision (step) that caused them, rather than being diluted across thousands of irrelevant tokens or lumped into a single trajectory score.

D. Systems Design

The paper outlines the necessary system architecture to support StepPO:

Step-Native Data: Data pipelines must store prompt IDs, response IDs, and rewards as structured steps.
Gateway & DataPool: A middleware layer (as seen in the Claw-R1 system) to ingest heterogeneous agents (white-box and black-box) and standardize them into step-native traces.
Asynchronous Training: Decoupling rollout generation from training to handle variable latency in agent interactions, utilizing prefix-tree merging to optimize compute for shared prefixes.

3. Key Contributions

Conceptual Framework: Proposes the shift from Token-Level MDP to Step-Level MDP as the correct abstraction for Agentic RL, arguing that the "step" is the true atomic unit of agent decision-making.
Algorithmic Innovation: Introduces StepPO, a PPO variant that performs credit assignment and policy updates at the interaction-step granularity, solving the misalignment between token-based generation and step-based rewards.
Data Representation: Defines a Structured Step-Level Representation that eliminates retokenization drift and enables precise context management.
Systems Architecture: Details the infrastructure requirements (Gateway, DataPool, Asynchronous Training) necessary to scale step-aligned learning, illustrated through the evolution of the authors' own projects (Agent-R1 and Claw-R1).
Empirical Validation: Provides experimental evidence showing StepPO's superiority over token-level baselines.

4. Experimental Results

Setup: The authors evaluated StepPO against a standard token-level PPO baseline on the HotpotQA dataset (a multi-hop reasoning task requiring multiple search steps).
Configuration: Both methods used the same base model (Qwen2.5-3B-Instruct), data, and rollout pipeline, differing only in the optimization algorithm (Step-level vs. Token-level).
Findings:
- Performance: StepPO consistently outperformed token-level PPO throughout the training process.
- Stability: The StepPO curve reached a higher peak score and maintained a stronger plateau in the middle-to-late training stages.
- Conclusion: Aligning the optimization unit with the interaction step provides a more effective learning signal for tasks requiring multi-step evidence collection and reasoning.

5. Significance and Impact

Paradigm Shift: The paper argues that Agentic RL is not merely an extension of LLM RL but requires a fundamental restructuring of the MDP, data representation, and credit assignment mechanisms.
Scalability: By addressing the system-level challenges (data heterogeneity, latency, compute efficiency), StepPO offers a practical path toward training robust, general-purpose agents.
Community Lens: It provides a unified "lens" for the community to understand agent behavior, suggesting that future advancements in agent training must move away from token-centric views toward step-aligned, interaction-aware frameworks.
Open Source: The authors reference open-source implementations (Agent-R1, Claw-R1) that embody these principles, facilitating further research and adoption.