Source-linked AI summary
Self-Improving Agents with World Model Predictive Paths
Haori Lu, Tianyu Gao, Yaoyao Liu, Yingying Li
TL;DR
Long-horizon language agents must make reliable decisions despite incomplete observations, and how future prediction and accumulated experience jointly support those decisions is underexplored. The paper combines world-model evaluation of candidate actions with MPPI-style cost weighting and a dual-branch experience memory. It reports the strongest overall performance across three backbones on ALFWorld and ScienceWorld, while noting increased inference latency and cost as the candidate set grows.
Problem
How future-outcome prediction and accumulated interaction experience can jointly support decision-making remains underexplored for agents facing long-horizon tasks and incomplete observations.
Source: §1 · p. 1Method
A test-time framework predicts outcomes for multiple candidate actions, evaluates their task progress and risk, applies MPPI-style weights, and uses two experience-memory branches to improve planning and next-observation prediction.
Source: §1 · p. 2, §4.1 · p. 4Results
Takeaways & Limitations
Measured computational costs remain practical for the evaluated tasks, while prediction and experience are used together to guide action selection.
Source: §1 · p. 2, §6 · p. 9Evaluating more candidates increases test-time world-model and evaluator calls, latency, and cost; the current adaptation uses retrieved textual experience with model parameters frozen.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 Introduction
Reliable decisions over long interactions remain challenging, and language environments lack explicit dynamics and costs. The paper combines world-model lookahead, cost-based action selection, and dual-branch experience memory; experiments report the strongest overall success-rate performance among evaluated methods across three backbones.
1 Introduction
- Long-horizon tasks require agents to sustain progress under incomplete observations, where poor decisions can cause repeated failures or wasted interactions.Source: §1 · p. 1
- Language environments lack explicit dynamics and costs, and better next-observation prediction alone does not necessarily improve task success.The paper therefore couples prediction directly with decision-making.Source: §1 · p. 1
- The framework predicts outcomes for several candidate actions, evaluates task progress and risk, and uses MPPI-style weights to guide the action ultimately selected or synthesized.Candidates are compared before being executed in the environment.Source: §1 · pp. 1–2
- A dual-branch experience memory improves candidate generation and action selection in the planner branch and next-observation prediction in the world-model branch.Both branches use retrieved textual context while model parameters remain frozen; the analysis bounds per-step regret using candidate coverage, action selection, and prediction errors.Source: §1 · p. 2
- Experiments on ALFWorld and ScienceWorld show the strongest overall success-rate performance among evaluated methods across three backbones.The paper also reports ablation studies and a computational cost analysis.Source: §1 · p. 2
2 Related Work
Prior work addresses interactive decisions through feedback, search, world-model prediction, or experience-based adaptation. This paper adapts MPPI-style prediction and weighting to language agents and uses interaction experience to improve both candidate selection and prediction.
2 Related Work
- LLM agents act in loops from textual or visual observations; ReAct (Yao et al., 2023b) and Say-Can ground actions in feedback or skills, ADaPT decomposes tasks, and other methods provide lookahead via branching or retrieved information.Source: §2 · p. 2
- The framework adapts MPPI’s sampled-candidate cost weighting to language agents, replacing continuous control updates with discrete action updates driven by language models.MPC executes the first action and replans after a new observation; MPPI weights sampled trajectories by predicted costs (Williams et al., 2015, 2017).Source: §2 · p. 2
- Unlike prior world-model approaches described here, the method keeps models frozen and uses interaction experience in two branches to improve prediction and candidate generation and selection.The related work contrasts this with methods using adaptive lookahead, structured knowledge, retrieved experience, or offline-built memory.Source: §2 · p. 2
- Where experience-based methods improve what agents propose, the dual-branch memory also improves what they predict, sharpening candidates and forecasts from one interaction stream.Source: §2 · pp. 2–3
3 Problem Formulation
The paper frames language-agent interaction as a long-horizon task with a goal, observations, history, and a finite interaction budget. It motivates a predict-and-replan decision loop using MPC and MPPI.
3 Problem Formulation
- The agent acts from a context containing the goal, prior interaction history, and current observation, with its policy defining a distribution over language actions.The interaction history records prior observations and actions.Source: §3 · p. 3
- The objective is to maximize the probability of completing the task within the interaction budget.Source: §3 · p. 3
- MPC selects the candidate action sequence with the lowest expected predicted cost, executes only its first action, and replans after the next observation.MPPI approximates this search by sampling candidate trajectories and favoring lower predicted costs with Gibbs weights.Source: §3 · p. 3
4 Method
The method predicts outcomes for candidate actions, scores them for task progress and risk, and uses cost-weighted guidance to choose an action. Dual-branch experience memory improves candidate generation and selection as well as future predictions.
4.1 MPPI-Inspired Decision Loop
- The proposal policy generates up to K atomic candidate actions that may advance the task, gather information, or recover from failures.The method uses one-step predictions rather than multi-step rollouts to simplify the architecture and control computational cost.Source: §4.1 · p. 3
- A world model predicts each candidate action’s immediate outcome, and an evaluator assigns bounded predictive costs for task progress and risk.The loop executes one action, updates the trajectory, and replans from the new context.Source: §4.1 · pp. 3–4, The objective has a unique minimizer · p. 4, Figure 1 · p. 4
The objective has a unique minimizer
- The evaluator favors outcomes with greater predicted goal usefulness, considering task progress, information value, and consistency with available evidence.The objective balances weighted predictive cost against a KL penalty from uniform weighting.Source: §4.1 · p. 4
- The minimizing Gibbs cost-tilted distribution weights candidates according to their cost differences scaled by λ; smaller λ emphasizes low-cost candidates, while larger λ retains more balanced influence.The main configuration uses λ = 10.Source: The objective has a unique minimizer · p. 4
4.2 Dual-Branch Experience Memory
- The regret decomposition identifies candidate coverage, predictive cost error, and controller selection as distinct factors addressed by the memory design.Candidate-generation memories target the coverage gap; world-model memories target predictive cost error.Source: §4.2 · p. 5
- The dual-branch memory uses planner experience to improve candidate generation and action selection, while world-model memory improves next-observation prediction.Both branches update textual memories while keeping model parameters frozen.Source: §4.2 · pp. 5–6
5 Experiments
Experiments evaluate the method against baselines on ALFWorld and ScienceWorld across three backbones, then examine action selection, ablations, and computational cost. The method ranks first in nine benchmark comparisons and second in the remaining three, while memory effects vary across benchmarks and inference is slower than for baselines.
5.1 Experimental Setup
- Evaluation covers 134 ALFWorld tasks and 90 ScienceWorld tasks, with trials limited to 30 environment steps and SR@1 and SR@5 as primary metrics.The baselines are ReAct, ITP-I, and WorldEvolver; the evaluation varies among Gemma-4-26B-A4B, Qwen3.5-9B, and GPT-5.4-mini, with policy and world-model parameters frozen during deployment.Source: §5.1 · p. 6
5.2 Results
- Across three backbones, Ours ranks first in nine comparisons and second in three; with Gemma-4-26B-A4B, it scores 65.7/84.3 on ALFWorld and 46.7/60.0 on ScienceWorld.Against ITP-I, the strongest baseline, Ours improves by 32.9/29.8 points on ALFWorld and 10.0/10.0 points on ScienceWorld.Source: §5.2 · p. 6
- 77.8–83.8% of decisions select a highest-weight candidate, while second-ranked candidates account for 10.4–12.2% and outside-proposal actions for 1.6–6.0%.The controller usually follows the highest-weight candidate but can select another candidate or generate a new action.Source: §5.2 · pp. 6–7, Figure 2 · p. 7
- In a ScienceWorld example, retrieved experience makes focusing a prerequisite for pickup; completing that subgoal lowers predicted pickup cost from 40 to 10 and precedes success.ReAct skips the focus subgoal and repeatedly declares completion, but fails.Source: §5.2 · p. 7, Figure 3 · p. 8
5.3 Ablation Study
- K = 8 yields the highest SR@1 on both benchmarks, exceeding K = 1 by 5.3 points on ALFWorld and 3.4 on ScienceWorld.At λ = 50, ALFWorld SR@1 falls from 65.7 to 53.7; performance is similar at λ = 1 and λ = 10, so λ = 10 is the default.Source: §5.3 · p. 7, Table 2 · p. 7, Table 3 · p. 7
- Enabling both memory branches produces the highest average SR@1 and SR@5, although their effects are not uniform across benchmarks.World-model memory alone raises ALFWorld SR@1 from 56.7 to 68.7 but hurts on ScienceWorld; planner memory alone helps on ScienceWorld but hurts on ALFWorld. Cumulative SR@1 curves initially coincide, then separate and remain apart on both benchmarks.Source: §5.3 · p. 7, Table 4 · p. 8, Figure 4 · p. 8
- At R = 4, increasing memory capacity C from 4 to 24 improves SR@1 on both benchmarks, but C = 36 lowers it; moderate budgets are relatively stable.R = 8 gives mixed results, and the shared default is (C, R) = (24, 4).Source: §5.3 · p. 7, Table 5 · p. 8
5.4 Computational Cost
- 6.72 seconds per step is slower than the baselines’ 0.51–1.94 seconds, with added overhead from predicting and scoring multiple candidate actions.Latency was measured on a task subset using Gemma-4-26B-A4B on 2× NVIDIA A40 GPUs; the authors consider the cost practical without strict real-time constraints when improved task success justifies it.Source: Table 6 · p. 8, §5.4 · p. 8
6 Conclusion
The framework uses world-model predictions and MPPI-style weights to guide decisions, while separate planner and world-model memories refine candidate selection and next-observation prediction. Experiments show strong overall performance across three backbones, with practical computational costs on the evaluated tasks.
6 Conclusion
- Experiments on ALFWorld and ScienceWorld show the strongest overall performance across three backbones, with computational costs remaining practical for the evaluated tasks.The framework predicts several candidate-action outcomes before acting and uses MPPI-style weights to summarize their costs.Source: §6 · pp. 8–9
- Planner memory refines candidate generation and selection, while world-model memory improves next-observation prediction.The two branches enable self-improvement from accumulated interaction experience.Source: §6 · p. 9
7 Limitations
The method’s adaptation relies on retrieved textual experience with model parameters frozen, and its multiple candidate evaluations add inference latency and cost. The experiments use text-based simulated environments; inaccurate predictions could lead to unsafe actions if deployed in real environments, where human oversight would be needed.
7 Limitations
- Adaptation uses retrieved textual experience while keeping model parameters frozen, and more capable adaptation to changing dynamics remains future work.World-model accuracy matters because predicted outcomes determine candidate costs and can change the controller’s action preference.Source: §7 · p. 9
- Evaluating multiple candidates adds world-model and evaluator calls at test time, increasing inference latency and cost as the candidate set grows.The authors identify more computationally efficient candidate comparisons as a future direction.Source: §7 · p. 9
- Experiments use text-based simulated environments; in real-world deployment, inaccurate world-model predictions could lead to unsafe actions, so human oversight would be needed.The experiments involve no human subjects or personal data.Source: §7 · p. 9
A World Model Memory Implementation Details
The world-model memory combines retrieved episodic transitions with reusable semantic rules to condition next-observation predictions. Both memory banks have bounded capacity and retrieve only a limited amount of experience for each prediction.
A World Model Memory Implementation Details
- For each candidate action, the episodic bank retrieves up to five transitions ranked by Jaccard similarity between tokenized candidate and stored actions.Each retrieved transition is presented as a state–action–next-state example alongside the task goal and current observation.Source: A World Model Memory Implementation… · p. 13
- The semantic bank extracts at most one reusable environment-change or persistence rule per prediction–observation mismatch and supplies up to eight recent rules to subsequent predictions.Source: A World Model Memory Implementation… · p. 13
- Across both benchmarks, episodic memory stores up to 100 transitions and semantic memory up to 128 rules, with memories retained across trajectories within each environment.Up to five transitions are retrieved per candidate and up to eight rules per prediction.Source: A World Model Memory Implementation… · p. 13
B Theoretical Analysis and Proofs
The proofs establish properties of the weighting rule and decompose per-step regret into candidate coverage, predictive cost error, and action selection. The supplied proof passages describe a unique minimizer via nonnegativity of a KL term and identify the controller’s selection gap.
B Theoretical Analysis and Proofs
- For a uniform distribution over candidates, the proof states that the resulting weights coincide with Eq. (10).Source: B Theoretical Analysis and Proofs · p. 13
- The proof establishes a unique minimizer by using nonnegativity of the KL divergence, which vanishes only at w = w⋆.It also states that subtracting the minimum yields the stated identity.Source: B Theoretical Analysis and Proofs · p. 13
- The regret decomposition includes candidate coverage, predictive cost error, and the final controller’s action-selection contribution.The selection gap captures the final controller’s contribution to per-step regret.Source: B Theoretical Analysis and Proofs · p. 13
C Detailed Results
Table 7 reports ALFWorld success rates by task type across three backbones; although the best method varies by task type, Ours has the best overall results for all three backbones.
C Detailed Results
- Ours achieves the best overall ALFWorld results across all three backbones, although the best method varies across task types.Table 7 reports SR@1 / SR@5 for PICK, CLEAN, HEAT, COOL, LOOK, and PICK2.Source: C Detailed Results · p. 14, Table 7 · p. 14
- The paper identifies overall success rate as the primary performance measure for the ALFWorld results.Source: C Detailed Results · p. 14
- Table 7 reports SR@1 / SR@5 for each of six ALFWorld task types: PICK, CLEAN, HEAT, COOL, LOOK, and PICK2.Each cell reports SR@1 / SR@5; bold and underlined overall success rates indicate the best and second-best results.Source: Table 7 · p. 14, C Detailed Results · p. 14
D Prompt Templates
The prompt templates revise a prior agent prompt for this framework and specify separate stages for proposing candidates, predicting outcomes, scoring them, selecting an action, and reflecting on feedback.
D Prompt Templates
- The framework substantially revises WorldEvolver’s ReactAgent prompt and provides module-specific prompts.Source: D Prompt Templates · p. 14
D.1 Candidate Generation
- Candidate generation proposes up to eight distinct one-step actions, grounded in visible names and preserving the goal’s requested operation.It can use information-gathering actions when they resolve relevant uncertainty, and should not repeat unchanged failures.Source: §D.1 · pp. 14–15
D.2 World Model Prediction
- The world model predicts the next observation from the current observation, proposed action, and goal, using retrieved transitions and frame rules as guidance.Source: §D.2 · p. 15
D.3 Cost Evaluation
- Cost evaluation jointly scores every candidate from 0 to 100, with lower costs indicating more likely progress or useful information.The scorer considers expected progress, information value, and consistency with visible evidence; forecasts are uncertain suggestions.Source: §D.3 · p. 16
D.4 Final Action Selection
- Final selection uses candidates, forecasts, and MPPI-inspired weights, while treating the current observation as authoritative and grounding action arguments in visible names.Source: §D.4 · p. 16
D.5 Planner Reflection
- Planner reflection labels observed transitions SUCCESS, FAILURE, or IGNORE, retaining experiences only for success and failure.Success and failure produce separate reusable experiences for candidate generation and final action selection; unresolved issues can also be recorded.Source: §D.5 · p. 17