Dependency aware reward shaping improves training for language agents
Dependency-Aware Reward Shaping for Agentic Reinforcement Learning
Artificial Intelligence
Summary
When computers learn to perform tasks by trial and error, they usually get rewards only when a task is completely done, which makes it hard to know which steps helped. The authors propose a new way to give the computer smaller rewards based on which sub-tasks have been completed or fixed during the process. This method uses a graph of task dependencies to more fairly assign credit for each step, leading to better training results. Their approach works with various existing training methods and improves performance on several complex language tasks.
What this means in practice
- •For language model trainers: Train large language models with better stepwise feedback for complex tasks using dependency-aware rewards to improve final success rates.
- •For ai tool developers: Integrate dependency-aware reward shaping into AI reasoning systems to enhance training efficiency without modifying their core algorithms.
Authors
Ziyi Chen, Yan Zhang, Jianhui Wei, Daoan Zhang, Zuozhu Liu
Abstract
When training large language models with reinforcement learning, terminal rewards provide little guidance about which steps matter. Common methods for assigning step credit overlook that work built on uncorrected mistakes is wasted while independent work remains valid. With only a final success/failure reward, every step in a failed episode has zero total future reward, even when it made progress. We propose Dependency-Aware Reward Shaping (DARS), which represents task progress as predicates linked by prerequisite relations and assigns step-level credit over the dependency graph. An annotator marks which predicates each step verifies, invalidates, or repairs. Verified predicates are discounted according to graph distance from the nearest broken prerequisite, while independent predicates are unaffected. Repairs update these weights based on any errors that remain; invalidated predicates need re-verification to regain credit. A fixed potential converts these annotations into signed per-step rewards. A common reward and annotation interface allows DARS to integrate with a range of reasoning and agentic training methods, such as GiGPO and ARPO/AEPO, without changing their rollout strategies or optimizers. Across five task families and models from 1.5B to 8B, DARS improves success by up to 10 points over GiGPO trained with the same budget and harness (ALFWorld), raises the WebShop task score and Search-R1 QA accuracy, complements AEPO's entropy-based training on AIME24/25 with a Python interpreter, and exceeds OmniOPD in controlled tool-free reasoning comparisons at 1.7B and 4B. Ablations show that step-level credit, dependency attenuation, and graph topology each contribute. On ALFWorld, a distilled 8B annotator matches the API annotator, enabling DARS to run efficiently without a frontier judge. Code is available at https://github.com/JianhuiWei7/DARS.