SALT: Step-level Advantage Assignment for Long-horizon Agents via Trajectory Graph

TL;DR

SALT is a lightweight, graph-based framework that assigns step-level advantages from outcome rewards to improve long-horizon RL tasks, enhancing performance without modifying existing algorithms.

Contribution

SALT introduces a novel graph-based advantage assignment method that can be integrated into group RL algorithms to better handle multi-step, long-horizon tasks.

Findings

Consistently improves performance across multiple benchmarks.

Seamlessly integrates with existing RL algorithms with minimal overhead.

Validates design choices through thorough analysis.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities, enabling language agents to excel at single-turn tasks. However, their application to complex, multi-step, and long-horizon tasks remains challenging. While reinforcement learning (RL) offers a promising avenue for addressing these challenges, mainstream approaches typically rely solely on sparse, outcome-based rewards, a limitation that becomes especially problematic for group-based RL algorithms lacking critic models, such as Group Relative Policy Optimization (GRPO). In such methods, uniformly rewarding or penalizing all actions within a trajectory can lead to training instability and suboptimal policies, because beneficial and detrimental actions are often entangled across multi-step interactions. To address this challenge, we propose SALT, a novel and lightweight framework that provides a finer-grained advantage assignment, derived solely from outcome rewards. We achieve this by constructing a graph from trajectories of the same prompt, which allows us to quantify the quality of each step and assign advantages accordingly. Crucially, SALT is designed as a plug-and-play module that seamlessly integrates with existing group-based RL algorithms, requiring no modifications to the rollout procedure and introducing negligible computational overhead. Extensive experiments on the WebShop, ALFWorld, and AppWorld benchmarks with various model sizes demonstrate that SALT consistently improves performance. We also conduct a thorough analysis to validate the design choices behind SALT and offer actionable insights.