Organizations: Johns Hopkins University, Baltimore MD, USA · Arizona State University, Tempe AZ, USA · University of Illinois Chicago, Chicago IL, USA · Purdue University, West Lafayette IN, USA
Large language models (LLMs) are increasingly deployed as multi-step decision-making agents, where effective reward design is essential for guiding learning. Although recent work explores various forms of reward shaping and step-level credit assignment, a key signal remains largely overlooked: the intrinsic uncertainty of LLMs. Uncertainty reflects model confidence, reveals where exploration is needed, and offers valuable learning cues even in failed trajectories. We introduce SELAUR: Self Evolving LLM Agent via Uncertainty-aware Rewards, a reinforcement learning framework that incorporates uncertainty directly into the reward design. SELAUR integrates entropy-, least-confidence-, and margin-based metrics into a combined token-level uncertainty estimate, providing dense confidence-aligned supervision, and employs a failure-aware reward reshaping mechanism that injects these uncertainty signals into step- and trajectory-level rewards to improve exploration efficiency and learning stability. Experiments on two benchmarks, ALFWorld and WebShop, show that our method consistently improves success rates over strong baselines. Ablation studies further demonstrate how uncertainty signals enhance exploration and robustness.
Figures & tables
Figure 1: Overview of our method, SELAUR. (a) Uncertainty Estimation : We employ three approaches to estimate model uncertainty: entropy, confidence, and margin, derived from token probability distributions. (b) Step and Trajectory Aggregation : Throughout the model’s process, we compute two types of rewards that contribute to the uncertainty reward. At the step level, token-level uncertainties are aggregated. At the trajectory level, step-level uncertainties are combined with varying weights, where later steps receive higher weights in the overall trajectory reward. (c) Failure-aware Reward Reshaping : For successful cases, the model is trained with the standard reward. For failed cases, we incorporate the uncertainty reward into the training loop to improve learning robustness.
Base Model
Method
AlfWorld
WebShop
Pick
Look
Clean
Heat
Cool
Pick2
All
Score
Succ.
Qwen2.5-1.5B
prompting
0.1071
0.0000
0.0000
0.071
0.0000
0.0000
0.0312
0.2180
0.0429
PPO
0.9142
0.3750
0.8333
0.8333
0.7368
0.7916
0.8046
0.7126
0.4921
RLOO
0.9000
0.7000
0.9200
0.7142
0.6875
0.5000
0.7812
0.8022
0.6394
GRPO
0.6756
0.5000
0.8000
0.9285
0.6190
0.5263
0.6796
0.8359
0.6718
GiGPO
0.9143
0.7500
1.0000
0.8889
0.8421
0.7916
0.8828
0.8017
0.6757
Table 1: Main experimental results on ALFWorld and WebShop . For ALFWorld , we report success rates (%) on individual tasks and their overall average (All). For WebShop , we report both the task score and success rate (%). Bold numbers mean the best performance, and italic numbers mean the second-best. Overall, our method achieves the strongest performance across both benchmarks.
Setting
Task Score
Success Rate
entropy
least confidence
margin
✗
✗
✗
0.8017
0.6757
✓
✗
✗
0.8504
0.7265
✗
✓
✗
0.8465
0.7226
✗
✗
✓
0.7834
0.6328
✓
✓
✗
0.8263
0.6953
Table 2: Each row evaluates a specific combination of metrics: ✓ indicates the inclusion of an uncertainty metric in the selection process, while ✗ denotes its exclusion. Results show that while individual or pairwise metrics provide partial gains, the full integration of Entropy, Least Confidence, and Margin ( ✓ across all) yields the optimal performance in both Task Score and Success Rate.
Figure 2: Entropy change during training under different uncertainty strategies. SELAUR maintains higher entropy throughout training, indicating stronger exploration compared to other methods.
Reward Method
AlfWorld
WebShop
Baseline
0.8828
0.6757
Negative
0.8516
0.7539
Exponential
0.8750
0.7031
Ours (SELAUR)
0.8906
0.7656
Table 3: Comparison of different reward functions on WebShop and AlfWorld . We report average of success rate(%) on AlfWorld and WebShop. Our proposed SELAUR achieves the best performance, demonstrating a better trade-off between exploration and exploitation.
Figure 3: Comparison of action traces in the WebShop environment. Top : SELAUR leverages uncertainty to explore alternative paths rather than being trapped in a single incorrect trajectory, eventually discovering the correct solution. Bottom : GiGPO exhibits low uncertainty, leading to repetitive and fixed behaviors.
Large language model (LLM)-based agents often make suboptimal tool-use decisions, including unsupported tool invocation and hallucinated direct responses, which may accumulate errors throughout multi-step interactions. Existing approaches mainly improve these behaviors through inference-time correction or coarse-grained reward signals based on decision outcomes and structured checklists, leaving the uncertainty characteristics of agent decisions underexplored. We observe that decision-oriented reinforcement learning tends to weaken the uncertainty separation between correct and incorrect actions, resulting in overconfident mistakes and weaker exploration signals. Therefore, we propose TRUST, which incorporates uncertainty quantification into reward design as a repulsive force for maintaining uncertainty separation, and labels lightweight key-turn annotations for unified post-training of multi-turn trajectories. Experimental results across diverse tool-use benchmarks show that TRUST consistently enhances both decision quality and agent performance while maintaining more reliable uncertainty estimates during optimization.
Yijin Zhou, Linqian Zeng, Xiaoya Lu +4
Shanghai Jiao Tong University, China · Shanghai Artificial Intelligence Laboratory, China · Shanghai Innovation Institute, China
Sparse rewards and heterogeneous task sequences remain persistent challenges in Reinforcement Learning (RL), often resulting in slow convergence, weak generalization, and inefficient exploration. We propose Uncertainty-Aware LLM-Guided Policy Shaping (ULPS), a novel framework that integrates a calibrated Large Language Model (LLM) into the RL training loop to provide structured, uncertainty-modulated behavioral guidance. ULPS employs an A*-based oracle to synthesize optimal symbolic trajectories, which are used to fine-tune a BERT-based language model. During training, this model supplies action suggestions whose influence is conditioned on epistemic uncertainty estimated via Monte Carlo (MC) dropout. An entropy-based blending mechanism adaptively balances LLM guidance and the learned policy (via Proximal Policy Optimization, PPO), allowing the agent to prioritize reliable priors while preserving adaptability. We evaluate ULPS on the MiniGridUnlockPickup benchmark and observe consistent improvements in success rate, reward efficiency, and sample complexity over unguided, uncalibrated, and standard RL baselines. ULPS achieves more than 9% improvement in execution accuracy after fine-tuning, requires fewer environment interactions, and yields higher reward AUC. Our results demonstrate that integrating symbolic A* trajectories, pretrained language priors, and uncertainty-aware control offers a principled and effective approach to multi-task reinforcement learning in sparse-reward domains, with potential extensibility to partially observable and multi-agent settings.
Reward models are central to aligning large language models (LLMs) with human preferences. Yet most approaches rely on pointwise reward estimates that overlook the epistemic uncertainty in reward models arising from limited human feedback. Recent work suggests that quantifying this uncertainty can reduce the costs of human annotation via uncertainty-guided active learning and mitigate reward overoptimization in LLM post-training. However, uncertainty-aware reward models have so far been adopted without thorough comparison, leaving them poorly understood. This work introduces a unified framework, RewardUQ, to systematically evaluate uncertainty quantification for reward models. We compare common methods along standard metrics measuring accuracy and calibration, and we propose a new ranking strategy incorporating both dimensions for a simplified comparison. Our experimental results suggest that model size and initialization have the most meaningful impact on performance, and most prior work could have benefited from alternative design choices. To foster the development and evaluation of new methods and aid the deployment in downstream applications, we release our open-source framework as a Python package. Our code is available at https://github.com/lasgroup/rewarduq.