Self-Improving Agents with World Model Predictive Paths

Haori Lu; Tianyu Gao; Yaoyao Liu; Yingying Li

cs.AIOct 20, 2026

Self-Improving Agents with World Model Predictive Paths

Authors: Haori Lu, Tianyu Gao, Yaoyao Liu, Yingying Li

Organizations: University of Illinois Urbana-Champaign

Abstract

Large language model (LLM) agents are increasingly expected to interact with environments and improve their decisions during deployment. However, how anticipating future outcomes and drawing on accumulated interaction experience can jointly support decision-making remains underexplored. We propose a test-time framework, inspired by Model Predictive Path Integral (MPPI) control, that combines multi-branch predictive analysis with experience-based adaptation. At each step, the agent proposes multiple candidate actions and imagines their outcomes with a world model, so that alternatives can be compared without being executed in the environment. Since language environments provide neither explicit dynamics nor costs, the predicted outcomes are evaluated for task progress and risk, and MPPI-style weights over the resulting costs guide the final decision. A dual-branch experience memory enables self-improvement: the planner branch refines candidate generation and action selection, while the world model branch improves next-observation prediction. Experiments on ALFWorld and ScienceWorld across three backbones demonstrate the strong overall performance of our method and the effectiveness of experience-guided predictive control. Our code will be made publicly available.

Figures & tables

Figure 1: Overview of our method. At step $t$ , the proposal policy takes the goal, history, and observation and proposes $K$ candidate actions $a_{t}^{(k)}$ . The world model predicts the next observation $\hat{o}_{t+1}^{(k)}$ of each candidate, the evaluator scores task progress and risk to give costs $C_{t}^{(k)}$ , and MPPI-style weights summarize these costs for the final controller, which decides the executed action $a_{t}$ . After execution, the transition updates both experience branches. The planner branch labels it as success, failure, or ignore and summarizes the first two into $M_{t}^{\mathrm{cand}}$ and $M_{t}^{\mathrm{final}}$ , retrieved by the proposal policy and final controller in future steps; the world model branch updates its episodic and semantic memory $M_{t}^{\mathrm{world}}$ , retrieved by the world model in future steps.

Panel of Table 1 — Table 1: Overall success rates (%) on ALFWorld and ScienceWorld across three backbones, reported as SR@1 and SR@5. Bold and underlined success rates indicate the best and second-best results, respectively.

Figure 2: Executed actions by candidate rank under MPPI-style weighting with Gemma-4-26B-A4B. Rank 1 includes ties for the highest weight; Outside denotes actions not in the proposal set.

Figure 3: An example of how retrieved experience and predictive evaluation guide action selection. Our agent completes the focus subgoal before pickup and succeeds, while ReAct skips this subgoal and fails.

$K$	ALFWorld		ScienceWorld		Avg.
$K$	SR@1	SR@5	SR@1	SR@5	SR@1	SR@5
1	60.4	82.1	43.3	58.9	51.9	70.5
2	59.7	85.1	38.9	56.7	49.3	70.9
4	61.9	82.8	40.0	58.9	51.0	70.9
8	65.7	84.3	46.7	60.0	56.2	72.2
16	64.2	84.3	42.2	68.9	53.2	76.6

Table 2: Effect of the number of candidate actions

K

imagined per step, with Gemma-4-26B-A4B (%). The default setting is shaded.

$\lambda$	ALFWorld		ScienceWorld		Avg.
$\lambda$	SR@1	SR@5	SR@1	SR@5	SR@1	SR@5
1	65.7	81.3	48.9	62.2	57.3	71.8
10	65.7	84.3	46.7	60.0	56.2	72.2
50	53.7	80.6	47.8	63.3	50.8	72.0

Table 3: Effect of weight temperature

\lambda

with Gemma-4-26B-A4B (%). Larger

\lambda

gives more uniform MPPI-style weights. The default setting is shaded.

Memory		ALFWorld		ScienceWorld		Avg.
Planner	WM	SR@1	SR@5	SR@1	SR@5	SR@1	SR@5
$\times$	$\times$	56.7	80.6	38.9	57.8	47.8	69.2
$\checkmark$	$\times$	49.3	73.9	44.4	60.0	46.9	67.0
$\times$	$\checkmark$	68.7	88.8	36.7	54.4	52.7	71.6
$\checkmark$	$\checkmark$	65.7	84.3	46.7	60.0	56.2	72.2

Table 4: Effect of planner and world model (WM) memory with Gemma-4-26B-A4B (%). ‘

\checkmark

’ indicates an enabled branch. The default setting is shaded.

Figure 4: Cumulative SR@1 over executed tasks with and without dual-branch memory. Shaded areas show the gain in percentage points.

$C$	$R$	ALFWorld		ScienceWorld
$C$	$R$	SR@1	SR@5	SR@1	SR@5
4	4	62.7	85.8	36.7	60.0
12	4	64.9	85.8	44.4	58.9
24	4	65.7	84.3	46.7	60.0
36	4	63.4	84.3	42.2	64.4
24	1	64.2	89.6	46.7	63.3

Table 5: Effect of planner memory capacity

C

and retrieval count

R

with Gemma-4-26B-A4B (%). The upper block varies

C

R=4

, the lower block varies

R

C=24

. The default setting is shaded.

Method	Avg. sec. / step $\downarrow$	SR@1 $\uparrow$	SR@5 $\uparrow$
ReAct	0.51	20.1	24.6
ITP-I	1.94	32.8	54.5
WorldEvolver	1.33	16.4	20.9
Ours	6.72	65.7	84.3

Table 6: Per-step latency and success rates using Gemma-4-26B-A4B on 2

\times

NVIDIA A40 GPUs.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Panel of Table 7 — Table 7: Success rates (%) on ALFWorld by task type across three backbones. Each cell reports SR@1 / SR@5. Bold and underlined overall success rates indicate the best and second-best results, respectively.

Self-Improving Agents

Self-improving agents use experience and feedback to update their models, memories, skills, or workflows for better future task performance.

See all 171 papers

Steering Black-Box Large Language Model Agents with Pitfall Anticipation World Model
Oct 21, 2026Jiongxiao Wang, Han Ding, Youzhi Luo +4LLM World Models LLM Agents
University of Wisconsin–Madison · AWS AI
Self-Evolving World Models for LLM Agent Planning
Jun 29, 2026Xuan Zhang, Wenxuan Zhang, See-Kiong Ng +1Test-Time Adaptation World Model Learning
National University of Singapore · Singapore University of Technology and Design · Singapore Management University
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

Oct 21, 2026J/K move · Enter open · S save

Steering Black-Box Large Language Model Agents with Pitfall Anticipation World Model

Jiongxiao Wang, Han Ding, Youzhi Luo, Yawei Wang+3

University of Wisconsin–Madison · AWS AI

LLM World Models LLM Agents Long-Horizon LLM Agents LLM Agent Reliability RL for LLM Agents

Open paper AskPDF arXiv

Frontier black-box Large Language Model (LLM) agents are proficient at long-horizon reasoning and precise action generation. However, when it comes to a specific agentic task, they have no internalized knowledge of how that particular environment unfolds. A world model offers a natural remedy, providing the agent with environment specific knowledge by foreseeing how the environment will respond to the action it is about to take. Prior world models for LLM agents mainly predict next observations at each step, which is costly and often uninformative because individual observations provide limited task-relevant guidance and not every step requires call for foresight. This raises a key question: what a world model should predict and when should it be invoked for agentic tasks? Our pilot study suggests that world model guidance is most useful when it anticipates pitfalls within a bounded multi-step future. Motivated by this, we introduce Pitfall Anticipation World Model (PAWM), which selectively intervenes when a potential pitfall arises and predicts the impending failure to steer the agent away from it. PAWM is trained in two stages: Supervised Fine-Tuning (SFT) on synthetic data labeled with agent pitfalls, followed by GRPO co-training with a black-box LLM to further optimize it in the target environment. PAWM improves Claude Sonnet 4.5 performance by 24.3% on ScienceWorld and 7.7% on WebArena-Lite, outperforming a range of baselines. These results demonstrate that PAWM provides an effective form of world model guidance for black-box LLM agents.

Abstract

Self-Improving Agents

Explore similar work

Steering Black-Box Large Language Model Agents with Pitfall Anticipation World Model

Self-Evolving World Models for LLM Agent Planning

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