Skip to content

Internalizing the Future: A Unified Agentic Training Paradigm for World Model Planning

๐Ÿ•’ Published (v1): 2026-06-25 19:05 UTC ยท Source: Arxiv ยท link

Why this paper was selected

Unified agentic training for world-model planning; addresses core LLM agent reactive failure in long-horizon tasks

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

๐Ÿ’ฌ Ask ChatGPTโœฆ Ask Claude

TL;DR

Standard LLM agents are reactive: they lack an internal world model to simulate future outcomes before committing to a plan. This paper proposes a three-stage training paradigm โ€” WM-AMT โ†’ FE-SFT โ†’ FC-RL โ€” that internalizes both a prospective state rollout and a verbalized \(Q\)-value estimate into a single autoregressive policy. The key insight is that post-training alone cannot inject genuine predictive grounding; capability must be built at the mid-training stage first.

Problem

Long-horizon LLM agents treat planning reactively: they condition only on past context, with no mechanism to evaluate "what happens if I follow this plan?" Attempts to inject look-ahead reasoning purely through SFT or RL on formatted foresight traces produce a format-capability gap โ€” agents mimic the syntactic structure of foresight but generate hallucinated, miscalibrated content that adds noise rather than signal. The root cause is that post-training can only elicit latent capabilities, not inject fundamentally new ones.

Method

A single causal LM is trained end-to-end through three sequential stages:

  1. WM-AMT (World Model Agentic Mid-Training): Large-scale agentic trajectory data \(\{ฯ„_i\}\) is augmented by inserting structured world model blocks \(z_t = [z_t^{(\text{traj})} \oplus z_t^{(\text{eval})}]\) at randomly sampled steps. \(z_t^{(\text{traj})}\) is an abstracted roadmap of anticipated future; \(z_t^{(\text{eval})}\) contains a semantic gap analysis plus a confidence score \(q_t \in [0, 100]\%\) (a verbalized analogue of \(Q(s,a)\)), estimated heuristically by an LLM from the ground-truth future trajectory \(\tau_{>t}\). The model is optimized via standard NTP loss over augmented trajectories on ~200B tokens, forcing autoregressive reconstruction of \(z_t\) conditioned strictly on history \(\leq s_t\).

  2. FE-SFT (Format-Eliciting SFT): A filtered, task-specific instruction-tuning dataset teaches the model to autonomously invoke and structurally delimit the world model block \(\hat{z}_t\) during inference. Because WM-AMT has already established predictive priors, hallucination risk is suppressed.

  3. FC-RL (Foresight-Conditioned RL): Using group-based RL (GRPO), the agent generates trajectories in an interleaved format: imaginary โ†’ keyword โ†’ analysis โ†’ confidence \((q_t)\) โ†’ think โ†’ action. Two reward streams are jointly optimized:

  4. World model rewards: (a) Grounding โ€” \(R_{\text{ground}}^{i,t} = R_{\text{match}}(K_t^{(i)}, \tau_{>t})\), rewarding keyword milestones that manifest in actual execution; (b) Calibration โ€” Brier score \(R_{\text{calib}}^{i,t} = I - (q_t - I)^2\), aligning \(q_t\) with terminal binary outcome \(I\).
  5. Policy rewards: Combined advantage \(A(\tau[\text{policy}]_{it}) = A_E(\tau_i) + \omega \cdot A_S(\tau_i^t)\), where the global episode advantage \(A_E\) handles sparse credit and the local step advantage \(A_S\) (derived from calibrated \(q_t\)) densifies it.

Key Contributions

  • Identifies and formalizes the format-capability gap: SFT/RL on look-ahead formats alone produces syntactically plausible but semantically hollow foresight.
  • Proposes WM-AMT, a mid-training stage that augments existing agentic trajectory corpora with compressed future experience summaries and verbalized \(Q\)-values, requiring no new human annotation.
  • Introduces FC-RL with dual reward streams (grounding + Brier-score calibration) that jointly refine the internal world model and the policy within a single model and no separate value head.
  • Demonstrates that a verbalized, calibrated confidence signal \(q_t\) outperforms state-only future prediction as a planning signal.

Results

Search tasks (7 QA datasets, evaluated with LLM-as-a-judge): - Full pipeline (WM-AMT + FE-SFT + FC-RL) achieves average 50.6, vs. 48.4 (IWM+RL), 48.7 (State-SFT+RL with mid-training), and 47.1 (FE-SFT+RL without WM-AMT) โ€” all on Youtu-LLM-2B. - WM-AMT alone lifts SFT-stage average from 38.2 โ†’ 39.9 (standard SFT) and from 37.4 โ†’ 41.8 (FE-SFT).

Mathematical reasoning (AIME 2024/2025/2026, 30 samples): - Full pipeline: mean@30 = 29.5, pass@30 = 60.0, vs. State-SFT+RLโ€ก (mean@30 = 27.7, pass@30 = 56.7) and FE-SFT+RL (mean@30 = not reported at best baseline level, pass@30 implicitly lower). - Gains are especially pronounced on multi-hop QA tasks (e.g., Bamboogle: 47.2 โ†’ 56.0).

Limitations

  • Experiments conducted on a single proprietary 2B model (Youtu-LLM-2B); generalizability to other architectures or scales is unverified.
  • The \(q_t\) confidence scores in WM-AMT are estimated heuristically by an LLM rather than via Monte Carlo rollouts, introducing noise in the initial training signal.
  • Evaluation domains are limited to search-augmented QA and mathematical reasoning; applicability to embodied or tool-use agents with richer action spaces is not demonstrated.
  • WM-AMT requires augmenting ~200B tokens of agentic trajectories, imposing significant compute cost that may be prohibitive outside well-resourced settings.
  • The calibration mechanism relies on a single terminal binary outcome \(I\) as a Monte Carlo return; this may be noisy for tasks without clear ground-truth correctness signals.

Relevance to Agentic AI / LLM Agents

This work directly addresses the central planning deficit of LLM agents in long-horizon tasks: the inability to perform look-ahead deliberation before committing to an action sequence. By internalizing world modeling into the policy itself (rather than relying on an external simulator), it sidesteps the simulator-policy mismatch and inference overhead that plague systems like WebDreamer. The three-stage capability-first pipeline is a practically significant training recipe finding โ€” it establishes that mid-training stage placement matters for injecting genuinely new cognitive capabilities, not just eliciting existing ones, which has broad implications for how future agent training curricula should be structured. The verbalized \(Q\)-value mechanism also provides a transparent, inspectable planning signal, relevant to interpretability and control in agentic systems.