AttnSR: Using Attention Mechanisms to Solve the Reward Assignment Challenge in Long-Horizon Agent Reinforcement Learning

Long-horizon LLM agents need to make multi-step decisions, but reinforcement learning that relies solely on outcome rewards struggles to provide direct supervision for each step. A new paper proposes the AttnSR method, which derives step-level reward assignment from the attention produced by the policy model during training rollouts, without training a separate reward model or requiring extra rollouts.

Core Problem: The Lack of Step-Level Supervision

In long-horizon agent tasks, the agent interacts with the environment for dozens or even hundreds of steps but ultimately receives only a single outcome reward signal. This sparse reward makes credit assignment difficult—the agent struggles to determine which intermediate steps are critical to final success. Traditional methods either rely on manually designed intermediate rewards, which generalize poorly, or train additional reward models, which is costly. AttnSR's starting point is straightforward: the policy model's own attention already implicitly encodes judgments about step importance—can this signal be used directly?

Credit Assignment is a classic challenge in reinforcement learning: when an agent undergoes a long sequence of actions and receives only a final outcome reward, it must determine which intermediate actions this reward should be attributed to. In long-horizon tasks, reward signals are sparse and delayed, causing the variance of policy gradient estimates to increase significantly—key steps and irrelevant steps in the same successful trajectory receive equal weight, diluting the learning signal. Common mitigation approaches include: designing dense intermediate rewards (relying on domain knowledge), training independent reward models (such as process reward models, PRMs, based on preference data), or employing Monte Carlo tree search-style multiple rollouts to estimate step values. These methods either have limited generalization or incur high computational and annotation costs. AttnSR specifically targets this cost-effectiveness trade-off.

How AttnSR Works

AttnSR's core idea is to trace back through the policy model's attention maps to estimate the correlation between earlier steps and the final response. Specifically, during training rollouts, the method traces backward along the policy model's attention map, converting the attention weights of the final response to each step into the basis for reward assignment. The terminal step's reward is measured against a running baseline, determining whether these assignments reinforce or penalize the selected steps. This design avoids training additional reward models or increasing environment interactions, making it a plug-and-play reward enhancement scheme.

You may not have noticed, but ablation experiments reveal a key insight: AttnSR's gains depend on how intermediate rewards are assigned, rather than simply accumulating step rewards over time. This means the attention-guided assignment mechanism itself is the source of performance improvement.

Running baseline is used in policy gradient algorithms to reduce the variance of reward signals. The approach maintains a reference value that updates as training progresses (such as a moving average of recent trajectory returns or a value network estimate), and subtracts this baseline from the actual terminal reward to obtain an advantage estimate. A positive advantage indicates the trajectory is better than average and should be reinforced; a negative one should be suppressed. AttnSR combines this mechanism with attention-based assignment: each step receives a share determined by attention weights, which is then multiplied by the advantage relative to the baseline. Thus the same assignment pattern can switch between reinforcement and penalty depending on the terminal outcome, avoiding dependence on absolute reward scales.

Training rollout refers to the process in which the agent interacts with the environment under the current policy, collects trajectories, and uses them for policy updates. AttnSR emphasizes that its attention signal is taken directly from this regular rollout phase rather than from additional evaluation rollouts, so it does not increase the number of environment interactions—which is precisely the key reason it achieves a 29.9% reduction in interactions on ScienceWorld.

Experimental Results: Dual Validation on ALFWorld and ScienceWorld

In a controlled study with Qwen2.5-7B-Instruct, AttnSR outperformed the PPO baseline at all evaluation checkpoints on ALFWorld, improving final held-out success rate from 91.2% to 97.1%. On ScienceWorld, AttnSR also showed improvement at every checkpoint while reducing environment interactions by 29.9%. The improvement in interaction efficiency is significant for real-world deployment—fewer environment calls mean lower training costs and faster iteration.

According to the paper, these results support using policy attention as a practical prior for step-level reward assignment in long-horizon agent reinforcement learning. The code is open-sourced on GitHub (fairyshine/AttnSR) for community reproduction and extension.

ALFWorld is a benchmark environment for embodied text agents. Tasks include finding and manipulating objects in simulated household scenes (e.g., "put a clean apple in the fridge"), requiring dozens of steps of exploration, localization, and manipulation—a common testbed for long-horizon credit assignment. ScienceWorld is an interactive science experiment environment where agents must complete experimental procedures involving measurement, heating, mixing, and other operations; these involve longer step sequences and higher demands on the correctness of intermediate decisions. Both benchmarks share the characteristics of sparse outcome rewards and semantically rich intermediate steps, making them suitable for evaluating step-level reward assignment methods. Held-out set success rate refers to the success rate on task instances not seen during training, measuring generalization rather than memorization.

Implications and Limitations

AttnSR's value lies in providing a lightweight credit assignment approach: it does not change the policy network architecture, introduces no additional models, and mines supervisory information solely from existing attention signals. This opens a new direction for optimizing training efficiency in long-horizon agent RL. However, the current experiments are based on a single model, Qwen2.5-7B-Instruct, and the method's generalization across different architectures and scales requires further validation. Additionally, the reliability of attention weights may degrade in longer sequence tasks, which is an issue that future research should address.