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; † and ⋆ denote in-domain and out-of-domain QA evaluation. QA reports EM (%); WebShop reports Success (%) and Score. Indented rows augment GRPO.
Alg.
Model
RL-only
w/ ECHO
w/ DGS-WM
Single-hop QA
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 κ with γ=1 fixed. Right: varying γ with κ=0.5 fixed. The gray dashed line marks the ECHO baseline (31.02%).
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.
Ning Lu, Baijiong Lin, Shengcai Liu +9
Southern University of Science and Technology · Hong Kong University of Science and Technology · Hong Kong University of Science and Technology (Guangzhou) +2
Recent studies on world modeling for Large Language Model (LLM) agents typically formulate the learning objective as next-observation prediction. However, this objective ties supervision to what a transition happens to reveal, which may omit the dynamics most relevant to the agent's current decision. To bridge this gap, we propose Agent-Authored World Modeling (AAWM), a training procedure that constructs supervision from the policy's own decision needs. Specifically, at each state, the agent identifies what it needs to understand about the environment before acting. These needs drive the retrieval of relevant transition evidence across trajectories, which is then synthesized into training targets that capture decision-oriented dynamics instead of reconstructing the next observation. This aligns the training objective with the dynamics the policy needs before acting, not with the contents of the next observation. Experimental results validate the effectiveness of AAWM across multiple environments and training settings. These results show that decision-aware world-model targets provide a more effective learning signal than next-observation prediction.
Large language models (LLMs) are being used as policies for autonomous decision-making and planning in many domains. Despite their strong reasoning capabilities, LLMs struggle with long-horizon tasks, especially under partial observability. World models are a promising way to enhance policy performance, both during training and inference. During inference, agents currently use world models to simulate the consequences of candidate actions before choosing an action, which can improve decision-making. However, we argue that simulation alone is an incomplete interface for decision-making under partial observability: simulation does not adequately capture uncertainty about the current state, which agents may need for accurate decision-making. We address this limitation with Belief-Based World Models (BB-WMs), which maintain a belief that LLMs can query to access information on what is known and uncertain about the current state. Before developing methods to learn accurate BB-WMs, this paper focuses on a more fundamental question: does exposing a world model's belief directly to an LLM policy improve decision-making? Our results show that giving LLM agents access to beliefs improves task performance under partial observability, while remaining complementary to existing simulation-based world models. Code: https://github.com/skumar-ml/belief-world-models.