StepOPSD: Step-Aware Online Preference Distillation for Agent Reinforcement Learning

2026-05-26

AI Summary

This work addresses the credit assignment mismatch between sparse trajectory rewards and critical local actions in multi-turn agent reinforcement learning. It proposes StepOPSD, a novel framework that refines credit assignment to the level of individual agent steps for the first time. StepOPSD guides GRPO policy updates through step-level trajectory decomposition, hindsight-augmented teacher context re-scoring, sign-preserving advantage shaping, and normalized credit budgeting. Additionally, it introduces a dual-parameter control mechanism—α_clip and λ_mix—to modulate learning dynamics. The method achieves state-of-the-art performance on ALFWorld (e.g., 79.1% on Heat, 95.0% on PickTwo) and Search-QA (61.6% on TriviaQA). Empirical analysis further reveals that α_clip stabilizes local trust regions, while λ_mix exhibits task-dependent tuning characteristics.

Abstract

Reinforcement learning for multi-turn agents suffers from a credit-assignment mismatch: rewards are sparse and trajectory-level, while success often hinges on a few local decisions. Existing online policy distillation (OPD) provides denser token-level supervision, but typically treats heterogeneous agent trajectories as monolithic strings rather than causal interaction units. We present StepOPSD, a post-rollout preference self-distillation framework that takes the agent step as the unit of credit redistribution. StepOPSD decomposes trajectories into action-centered step segments, rescoring them under hindsight-enriched teacher contexts and converting token-level log-probability gaps into sign-preserving advantage shaping with a normalized per-step credit budget before the GRPO update. Across ALFWorld and Search-QA with Qwen3-1.7B and Qwen2.5-3B-Instruct, StepOPSD attains best or second-best results on subsets most sensitive to local causal errors, including first-place performance on ALFWorld Heat (79.1%), PickTwo (95.0%), Search-QA TriviaQA (61.6%), and tied-best performance on HotpotQA (40.4%). The results further reveal a consistent two-knob law: smaller α_clip acts as a broadly stabilizing local trust region, whereas the optimal global mixing strength λ_mix remains task-dependent. These findings suggest that step-aware distillation is most useful when trajectory-level rewards are weakly aligned with the local action that determines downstream success.

Problem

• credit assignment
• reinforcement learning
• multi-turn agents
• sparse rewards
• trajectory-level rewards

Innovation

• step-aware distillation
• online preference distillation
• credit assignment
• reinforcement learning
• advantage shaping