Decision aligned training improves long horizon reinforcement learning agents
Beyond Timestamps: Decision-Aligned On-Policy Distillation for Long-Horizon Agents
Machine LearningArtificial Intelligence
Summary
Reinforcement learning agents often learn from sparse rewards that come only at the end of a long task, making training difficult. The authors point out that existing training methods can misalign feedback with the agent’s actual decisions, especially when decisions span multiple time steps. They propose a new method called AlignOPSD that better matches feedback to the functionally correct parts of the agent’s behavior and distributes credit accordingly. This approach improves learning performance across multiple tasks compared to previous methods.
What this means in practice
- •For robotics engineers: Train robots to handle complex tasks requiring decisions over long time periods with improved accuracy from aligned feedback.
- •For ai product developers: Develop AI assistants that perform multi-step online tasks like web shopping and search more reliably through better reward crediting.$Commercial implications: Enables creation of commercial AI systems that learn complex multi-stage tasks more efficiently, improving product quality.
Authors
Mingju Chen, Can Lv, Jinrong Liu, Huan Zhang, Heng Chang, Shiji Zhou
Abstract
Reinforcement learning with verifiable rewards (RLVR) often relies on sparse outcome rewards, providing coarse supervision for long-horizon agents. On-policy self-distillation (OPSD) complements this signal with dense privileged feedback. However, we identify \emph{Decision--Timestamp Mismatch}: privileged guidance may be misaligned with the student's functional decision because the corresponding decision can occur at a different timestep, while the student's decision itself may span multiple timesteps rather than being tied to a single timestamp. Thus, timestamp-local supervision can misalign both the context and the temporal scope of credit. To address this mismatch, we introduce \textsc{AlignOPSD}, following the principle of aligning supervision before assigning credit. Decision-Aligned Supervision Rectification re-scores the same student-sampled response in functionally matched contexts across sibling rollouts to calibrate local teacher evidence. Semi-Markov Hierarchical Credit Assignment then derives variable-duration decision spans from correspondence changes and uses rectified evidence to allocate outcome-grounded credit across spans and their constituent turns. We evaluate \textsc{AlignOPSD} with Qwen2.5-3B and Qwen2.5-7B on ALFWorld, WebShop, and Search-QA against representative baselines. \textsc{AlignOPSD} outperforms both GRPO and StepOPSD across all eight backbone--aggregate-metric comparisons, improving on GRPO by 5.5--8.7 \% and ranking first in six. Additional analyzes examine the two alignment stages and hyperparameter sensitivity between tasks. Our code is avaliable at https://github.com/mingju-c/Align-OPSD