Model What Matters: Decision-Guided Selective World Modeling for Language Agents

Maolin Ran; Minda Chen; Di Huang; Wei Xiong; Fei Huang; Weiwen Liu

cs.AIOct 14, 2026

Model What Matters: Decision-Guided Selective World Modeling for Language Agents

Authors: Maolin Ran, Minda Chen, Di Huang, Wei Xiong, Fei Huang, Weiwen Liu

Organizations: Shanghai Jiao Tong University · LongShine AI Research

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.

Figures & tables

Figure 1: Comparison of typical policy-world modeling co-training versus DGS-WM on a two-hop question answering task. Given the same environment observations, typical co-training methods supervise all observation content equally as WM targets (left), while DGS-WM selectively supervises only decision-relevant content based on their support for future policy responses (right). Green highlights indicate selected content; gray text indicates filtered noise.

Figure 2: Overview of DGS-WM. Complete trajectories provide RL signal; observations and subsequent responses are used for target selection; history, action, and selected targets are used for world modeling prediction, producing target-level confidence and weights; weighted world modeling loss jointly updates policy with RL loss.

		Single-hop QA				Multi-hop QA					WebShop
Model	Method	NQ †	Trivia ⋆	PopQA ⋆	Avg	HotpotQA †	2Wiki ⋆	MusiQ ⋆	Bamb ⋆	Avg	Succ.	Score
Qwen3- 1.7B	Direct	16.18	33.99	22.12	24.10	11.03	3.75	2.03	3.20	5.00	16.60	45.51
	ReAct	18.20	36.52	22.43	25.72	12.98	11.07	4.01	20.00	12.02	17.20	54.70
	GRPO	36.29	55.41	42.22	44.64	26.46	24.20	5.67	19.20	18.88	37.80	64.21
	w/ ECHO	37.09	56.71	39.16	44.32	27.08	21.06	7.49	15.20	17.71	40.40	67.93
	w/ DGS-WM	37.53	57.46	42.31	45.77	30.30	25.45	7.70	24.80	22.06	43.20	70.39

Table 1: Main results. QA models train on NQ and HotpotQA;

\dagger

and

\star

denote in-domain and out-of-domain QA evaluation. QA reports EM (%); WebShop reports Success (%) and Score. Indented rows augment GRPO.

Single-hop QA
Alg.	Model	RL-only	w/ ECHO	w/ DGS-WM
GRPO	Qwen2.5-1.5B	45.34	44.64	46.28
GRPO	Llama-3.2-3B	47.87	49.82	50.75
GRPO	Qwen2.5-7B	49.76	50.55	51.99
PPO	Qwen2.5-1.5B	44.15	43.68	46.06
RLOO	Qwen2.5-1.5B	44.15	44.97	45.85

Table 2: Generalization across models and RL algorithms. EM (%) for single-hop and multi-hop QA averaged over all benchmarks in each category. Each row’s “RL-only” result uses the listed algorithm alone; “w/ ECHO” and “w/ DGS-WM” are its corresponding augmented variants.

Method	Single-hop QA	Multi-hop QA
GRPO	44.64	18.88
GRPO w/ ECHO	44.32	17.71
GRPO w/ DGS-WM	45.77	22.06
w/o DGOS	44.45	20.31
w/o CAW	44.82	21.66

Table 3: Ablation results on Qwen3-1.7B (EM, % by QA category). The first three rows compare GRPO with its ECHO- and DGS-WM-augmented variants; the final two rows ablate the indicated DGS-WM components.

Figure 3: Parameter sensitivity on Qwen3-1.7B. Combined QA EM (%) is the equally weighted average of single-hop and multi-hop averages. Left: varying

\kappa

with

\gamma=1

fixed. Right: varying

\gamma

with

\kappa=0.5

fixed. The gray dashed line marks the ECHO baseline (31.02%).

World Models

World models represent and predict how environments change, helping agents plan, act, and reason about possible futures.

See all 369 papers

Policy and World Modeling Co-Training for Language Agents
Jun 1, 2026Ning Lu, Baijiong Lin, Shengcai Liu +9Agentic RL World Model Learning
Southern University of Science and Technology · Hong Kong University of Science and Technology · Hong Kong University of Science and Technology (Guangzhou) +2
Beyond Next-Observation Prediction: Agent-Authored World Modeling for Sequential Decision Making
Jun 24, 2026Guangfeng Cai, Kaibing Yang, Shuo He +4World Model Learning World Models
Southeast University · Meituan
Towards a Belief-Based World Model for LLM Agents
Aug 31, 2026Shubham Kumar, Harshit Kumar, Narendra Ahuja +1Belief-Space Planning Partially Observable RL
UIUC · IBM

Jun 1, 2026J/K move · Enter open · S save

Policy and World Modeling Co-Training for Language Agents

Ning Lu, Baijiong Lin, Shengcai Liu, Jiahao Wu+8

Southern University of Science and Technology · Hong Kong University of Science and Technology · Hong Kong University of Science and Technology (Guangzhou)+2

Agentic RL World Model Learning LLM Agent Training RL for LLM Agents

Open paper AskPDF arXiv

Reinforcement learning (RL) improves large language model (LLM) agents by teaching them which actions lead to high rewards, but provides little supervision on what those actions do to the environment. World modeling (WM) can fill this gap, yet existing approaches often require separate simulators, extra training stages, or additional inference-time computation. We observe that on-policy RL rollouts already contain the needed signal: each transition pairs an action with its resulting next observation. Based on this observation, we propose PaW, a Policy and World modeling co-training framework that adds auxiliary WM supervision to the same policy during RL, without changing the inference paradigm. To make auxiliary WM supervision informative and stable, PaW introduces three components: action-entropy-based WM data selection, noise-tolerant WM loss, and reward-adaptive loss balancing. Experiments on three agentic task benchmarks show consistent improvements over strong RL baselines across models and RL algorithms. These results suggest that standard RL rollouts are a practical source of WM supervision for language-agent training.

Abstract

World Models

Explore similar work

Policy and World Modeling Co-Training for Language Agents

Beyond Next-Observation Prediction: Agent-Authored World Modeling for Sequential Decision Making

Towards a Belief-Based World Model for LLM Agents