ARIA: Training Language Agents with Intention-Driven Reward Aggregation

2025-05-31

AI Summary

In open-ended language-action environments, the exponential growth of the action space leads to sparse rewards and high variance in policy gradients. To address this, we propose Intention-Space Reward Aggregation (ISRA): a method that maps natural-language actions into a low-dimensional semantic intention space and enables cross-action reward sharing via semantic clustering—thereby densifying reward signals and stabilizing policy optimization. ISRA introduces, for the first time, a learnable intention encoder and a semantic-similarity-driven reward reweighting mechanism, making it compatible with mainstream LLM-based policies and both online and offline RL paradigms. Evaluated across four benchmark tasks, ISRA achieves an average performance improvement of 9.95% over strong baselines while significantly reducing policy gradient variance, consistently outperforming all comparative approaches.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Intelligent Robots: Learning & Optimization for ROB
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Large language models (LLMs) have enabled agents to perform complex reasoning and decision-making through free-form language interactions. However, in open-ended language action environments (e.g., negotiation or question-asking games), the action space can be formulated as a joint distribution over tokens, resulting in an exponentially large action space. Sampling actions in such a space can lead to extreme reward sparsity, which brings large reward variance, hindering effective reinforcement learning (RL). To address this, we propose ARIA, a method that Aggregates Rewards in Intention space to enable efficient and effective language Agents training. ARIA aims to project natural language actions from the high-dimensional joint token distribution space into a low-dimensional intention space, where semantically similar actions are clustered and assigned shared rewards. This intention-aware reward aggregation reduces reward variance by densifying reward signals, fostering better policy optimization. Extensive experiments demonstrate that ARIA not only significantly reduces policy gradient variance, but also delivers substantial performance gains of an average of 9.95% across four downstream tasks, consistently outperforming offline and online RL baselines.

Problem

• Addresses extreme reward sparsity in large language action spaces
• Reduces reward variance by clustering semantically similar actions
• Improves policy optimization in open-ended language environments

Innovation

• Projects actions into low-dimensional intention space
• Clusters semantically similar actions for shared rewards
• Reduces reward variance to optimize policy effectively