Skip to content

DreamPhase: Offline Imagination and Uncertainty-Guided Planning for Large-Language-Model Agents

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; DreamPhase: offline imagination and uncertainty-guided planning for agents

Ask a follow-up

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

💬 Ask ChatGPT✦ Ask Claude

TL;DR

DreamPhase is a modular planning framework that lets a frozen LLM agent simulate multi-step futures in a learned latent space before acting, scoring imagined branches by value minus an uncertainty penalty and distilling the best branch into a natural-language reflection that conditions the next prompt. This decouples environment-dynamics modeling from policy reasoning, achieving strong task performance while reducing real environment interactions by ~4× and irreversible actions by ~5× compared to online search baselines.

Problem

LLM agents face a safety-efficiency tension: online search planners (ReAct+MCTS) require many real environment queries and risk executing irreversible actions, while greedy imitation agents lack foresight and derail on early mistakes. Neither approach can safely plan ahead without costly or dangerous real-world interaction, and closed API LLMs cannot be fine-tuned on task-specific data.

Method

DreamPhase trains a latent world model (LWM) offline on logged trajectories \((ι, x_t, a_t, x_{t+1})\) collected from a frozen LLaMA-2-7B ReAct policy. The LWM has three components: an encoder \(f_θ(x_t) = z_t\), a stochastic transition model \(z_{t+1} \sim q_θ(z_{t+1}|z_t, \bar{a}_t, ι)\), and a decoder \(\hat{x}_{t+1} = d_θ(z_{t+1}, ι)\), trained with cross-entropy reconstruction plus KL regularization toward \(\mathcal{N}(0,I)\).

At each timestep, the agent rolls out \(M\) latent branches of horizon \(H\) entirely offline. Each branch \(\tau^{(j)}\) is scored by a discounted value head \(G^{(j)} = \sum_{k=1}^H \gamma^{k-1} V_\phi(z_{t+k}^{(j)}|ι)\) and penalized by an epistemic uncertainty estimate \(u^{(j)}\) computed via mutual-information proxy with Monte Carlo dropout on the frozen policy:

\[u^{(j)} = \mathcal{H}[\bar{p}^{(j)}] - \frac{1}{N}\sum_{n=1}^N \mathcal{H}[p^{(j)}(\xi_n)]\]

The penalized score is \(\tilde{G}^{(j)} = G^{(j)} - \beta u^{(j)}\). A safety gate passes the best branch \(j^*\) only if \(u^{(j^*)} \leq \tau\); otherwise the agent falls back to querying the real environment. The selected branch is distilled into a short natural-language reflection \(c_t\) and trajectory summary \(s_t\) (both within a 30-token budget) via lightweight frozen heads \(R_\phi\) and \(S_\eta\), which are prepended to the frozen policy LLM's prompt for the next action.

Key Contributions

  • Latent imagination planning: offline \(M\)-branch, \(H\)-step rollouts in latent space using a learned stochastic world model, avoiding environment queries during planning.
  • Uncertainty-aware branch selection: risk-sensitive score \(\tilde{G}^{(j)} = G^{(j)} - \beta u^{(j)}\) plus a confidence safety gate with theoretical regret bound \(\mathcal{O}(\sqrt{T\varepsilon}) + B\rho T\) linking performance to model error \(\varepsilon\) and mis-gating rate \(\rho\).
  • Zero-tuning language steering: reflections and trajectory summaries injected as prompt context; the policy LLM is never fine-tuned.
  • Sample and cost efficiency: <10 API calls/episode on WebShop vs. ~40 for ARMAP-M; ~12 ms imagination overhead; ~5× reduction in executed irreversible actions on WebShop, 4.9× on ALFWorld.

Results

  • Table 1 (8 diverse tasks, LLaMA-2-7B backbone): DreamPhase avg 50.1 vs. ARMAP-M 42.3, AgentGym 39.1, AgentLM 5.3; matches or exceeds closed-source GPT-4-Turbo (avg 63.9 across tasks where comparisons are shown).
  • SciWorld: DreamPhase 72.4 vs. ARMAP-M 51.2; TODOList: 45 vs. 35; TextCraft: 34 vs. 17.
  • Table 2 (WebShop + SciWorld + Game-of-24, cross-backbone):
  • LLaMA-70B: DreamPhase avg 53.6 vs. ARMAP-M 51.2.
  • LLaMA-8B: DreamPhase avg 35.1 vs. ARMAP-M 31.7.
  • Mistral-7B: DreamPhase avg 30.7 vs. ARMAP-M 28.9.
  • Phi-3.8B: DreamPhase avg 31.4 vs. ARMAP-M 29.1.
  • Game-of-24 gains are consistent across backbones (+4–9 points over ARMAP-M).
  • WebShop API calls: <10 per episode (DreamPhase) vs. ~40 (ARMAP-M).

Limitations

  • The latent world model requires offline logged trajectories from each benchmark's training split; distribution shift between training and test environments degrades accuracy and triggers more safety-gate fallbacks.
  • The regret bound relies on a uniform KL bound \(\varepsilon\) and Lipschitz value function assumptions that may not hold in complex, highly dynamic environments.
  • Imagination quality depends on the quality of the initial logged trajectories (collected by a simple ReAct LLaMA-2-7B); environments with sparse or poorly structured observations (e.g., visual frames rather than DOM trees) may yield noisier world models.
  • Monte Carlo dropout as a mutual-information proxy is an approximation; miscalibration can inflate or deflate \(u^{(j)}\), affecting gate reliability.
  • Evaluation is limited to benchmarks (WebShop, SciWorld, BabyAI, etc.); real-world deployment with live irreversible consequences (e.g., actual e-commerce) is untested.

Relevance to Agentic AI / LLM Agents

DreamPhase directly addresses the core open problem of making LLM agents safe and sample-efficient in multi-step decision-making without requiring fine-tuning or expensive online search. The approach is modular — it wraps any frozen LLM policy — making it practically applicable to closed-API models like GPT-4 and Claude that dominate real deployments. The combination of world-model imagination, uncertainty-gated action selection, and language-based reflection connects three active threads in agentic AI: model-based planning (Dreamer-style latent dynamics), uncertainty quantification for safe exploration, and self-reflection prompting (Reflexion). The regret bound provides a rare theoretical grounding for LLM agent planning that the field generally lacks.