SELAUR: Self Evolving LLM Agent via Uncertainty-aware Rewards

1. Problem Statement

Large Language Models (LLMs) are increasingly deployed as multi-step decision-making agents. While Reinforcement Learning (RL) is a standard paradigm for optimizing these agents, existing approaches suffer from two critical limitations:

• Sparse Reward Signals: Conventional RL often relies solely on binary outcome rewards (success/failure) at the end of a trajectory. This provides limited guidance for intermediate reasoning steps, making it difficult for agents to learn from failed attempts.
• Neglect of Intrinsic Uncertainty: Most step-level credit assignment methods ignore the model's intrinsic uncertainty (confidence). They fail to utilize the fact that high uncertainty indicates a need for exploration, while low uncertainty indicates reliable behavior. Consequently, agents often get trapped in local optima or fail to extract learning signals from failed trajectories where specific steps might have been promising but the final outcome was incorrect.

2. Methodology: SELAUR

The authors propose SELAUR (Self Evolving LLM Agent via Uncertainty-aware Rewards), a reinforcement learning framework that integrates token-level uncertainty directly into the reward design. The framework consists of three core modules:

A. Uncertainty Estimation

Instead of relying on a single metric, SELAUR aggregates three complementary token-level uncertainty measures to capture different facets of predictive reliability:

1. Entropy: Measures the dispersion of the probability distribution over the vocabulary. High entropy implies the model is unsure among many tokens.
2. Least Confidence: Based on the probability of the selected token ($1 - P_{selected}$). Lower probability indicates higher uncertainty.
3. Margin: Measures the gap between the top two token probabilities. A small margin indicates ambiguity between the top choices.

These metrics are combined into a unified token-level uncertainty score ($u_{t,j}$) using tunable weights.

B. Hierarchical Aggregation

To make token-level signals usable for RL, they are aggregated at two levels:

• Step-level: Token uncertainties within a single action step are averaged to produce a step-level uncertainty ($u_{step}$).
• Trajectory-level: Step-level uncertainties are aggregated across the entire trajectory ($\tau$) using an exponentially discounted mechanism. This assigns higher weights to later steps, as they often determine the final success or failure of the task.

C. Failure-aware Reward Reshaping

This is the core innovation. SELAUR modifies the standard reward function to incorporate uncertainty signals, specifically transforming failed trajectories into learning opportunities:

• Successful Trajectories: The agent receives standard rewards.
• Failed Trajectories: Instead of discarding the trajectory or assigning a flat negative reward, the framework injects the uncertainty signal:
  • Step-wise Shaping: Each step in a failed trajectory receives an augmented reward based on its normalized uncertainty. This encourages the agent to revisit steps where it was uncertain, treating them as valuable exploration points rather than just "failures."
  • Trajectory-level Shaping: The overall trajectory reward is adjusted based on the aggregated trajectory uncertainty.

3. Key Contributions

1. Unified Uncertainty Framework: The paper introduces a method to integrate multiple token-level uncertainty metrics (Entropy, Least Confidence, Margin) to create a dense, confidence-aligned supervision signal.
2. Failure-aware Reward Reshaping: A novel mechanism that uses uncertainty signals to reshape rewards for failed trajectories. This allows the agent to extract meaningful learning signals from unsuccessful attempts, improving exploration efficiency and policy stability.
3. Empirical Validation: Extensive experiments demonstrate that modeling failure signals via uncertainty leads to more robust policy learning compared to strong baselines.

4. Experimental Results

The method was evaluated on two interactive benchmarks: ALFWorld (instruction following in household tasks) and WebShop (goal-oriented e-commerce decision making).

• Performance: SELAUR consistently outperformed strong baselines (PPO, RLOO, GRPO, and GiGPO) across different model sizes (Qwen2.5-1.5B and 7B).
  • On ALFWorld, SELAUR achieved the highest average success rate (e.g., 96.87% with the 7B model vs. 94.53% for the next best).
  • On WebShop, it achieved the highest Task Score (0.8935) and Success Rate (79.68%).
• Ablation Studies:
  • Metric Combination: Using all three uncertainty metrics (Entropy + Least Confidence + Margin) yielded the best performance, proving their orthogonality and complementarity.
  • Reward Design: The proposed "Failure-aware" reshaping outperformed variants that simply penalized uncertainty (Negative variant) or discounted successful rewards based on uncertainty (Exponential variant).
• Qualitative Analysis: Case studies showed that SELAUR agents avoid "uncertainty collapse" (repetitive loops seen in baselines like GiGPO) by dynamically adjusting exploration based on confidence levels, eventually escaping local optima to find correct solutions.

5. Significance

SELAUR addresses a fundamental gap in LLM agent training by bridging the divide between sparse outcome-based rewards and dense token-level learning signals.

• Efficiency: It improves sample efficiency by leveraging information from failed trajectories that traditional RL would discard.
• Robustness: By explicitly modeling uncertainty, the agent learns to balance exploitation (using confident knowledge) and exploration (investigating uncertain states), leading to more stable policy convergence.
• Generalizability: The framework is model-agnostic and can be applied to various interactive environments, offering a new paradigm for self-evolving agents that learn effectively from both success and failure.