SELAUR: Self Evolving LLM Agent via Uncertainty-aware Rewards
2026-02-24 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors focus on improving how large language models (LLMs) make decisions step-by-step by considering how uncertain the models are about their choices. They developed a method called SELAUR that uses different ways to measure uncertainty and includes these signals in the reward system during training. This helps the model explore better and learn more stably, even when it fails sometimes. Tests on two tasks showed that SELAUR helped LLM agents succeed more often than other methods. Their additional experiments confirmed that incorporating uncertainty helps the model be more robust and explore more effectively.
large language modelsreinforcement learningreward shapinguncertainty estimationentropyleast-confidencemarginexplorationcredit assignmentALFWorld
Authors
Dengjia Zhang, Xiaoou Liu, Lu Cheng, Yaqing Wang, Kenton Murray, Hua Wei
Abstract
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.