Decision guided world modeling improves language agent task learning
Model What Matters: Decision-Guided Selective World Modeling for Language Agents
Artificial IntelligenceComputation and LanguageMachine Learning
Summary
Language agents learn to perform tasks by trying actions and seeing if they succeed, but they rarely get feedback about what happens after each action. This limits their understanding of how their actions change the world. The authors propose a method that focuses the learning on parts of the environment most relevant for making decisions, rather than everything equally. Their method also adjusts how strongly it learns from different pieces of information based on how confident the model is. This approach helps language agents perform better and more consistently on tasks like answering questions and interacting with websites.
What this means in practice
- •For software developers: Improve language-based assistants by training agents that better understand the impact of their actions in complex environments.
- •For automated customer support teams: Create more reliable conversational agents that selectively focus on important user inputs to guide responses in multi-step interactions.
Authors
Maolin Ran, Minda Chen, Di Huang, Wei Xiong, Fei Huang, Weiwen Liu
Abstract
Reinforcement learning provides a direct mechanism for optimizing task objectives in large language model-driven agents, but task rewards are typically sparse trajectory-level signals that only evaluate task success, without directly supervising action-conditioned environment transitions, leaving the model with limited explicit understanding of action consequences and environment dynamics. Recent work introduces auxiliary supervision through policy and world modeling co-training, yet typically treats complete environmental observations uniformly as world modeling targets, overlooking the varying decision relevance of observation content. We propose DGS-WM (Decision-Guided Selective World Modeling), a decision-guided policy-world modeling co-training approach for language agents. DGS-WM employs a Decision-Guided Observation Selection component that filters observation units with decision value as world modeling targets based on their support for subsequent policy responses, and a Confidence-Adaptive Weighting component that adaptively adjusts supervision strength according to the model's prediction confidence on selected targets. Together, these components form a decision-aware auxiliary world modeling loss, jointly optimized with the RL loss on the same policy model. Across multiple benchmarks covering single-hop/multi-hop question answering and interactive web tasks, DGS-WM achieves consistent improvements under different model scales and RL algorithms, demonstrating more stable and effective performance compared to full-observation world modeling methods.