Second Thoughts: Exploring and Planning through Comparative Imagination in World Models

Hieu Tran; Hong Yu

cs.LGOct 20, 2026

Second Thoughts: Exploring and Planning through Comparative Imagination in World Models

Authors: Hieu Tran, Hong Yu

Organizations: University of Massachusetts Amherst · University of Massachusetts Lowell

Abstract

World models let agents compare imagined outcomes before acting, but an agent must still decide which alternatives to explore and when a proposed plan should replace its learned policy. We introduce SDI (Sibling Disagreement in Imagination), which connects these decisions through a shared actor-based reference. SDI compares action sequences imagined from the same model state, using repeated rollouts of the actor's own sequence to measure variation under fixed behavior. This reference scales both the spread across candidate returns, which provides an exploration signal, and a proposed plan's predicted advantage, which governs action acceptance. The agent executes a planned action when its advantage exceeds the reference variability by a fixed factor; otherwise, it follows the actor. SDI reuses the base agent's world model, actor, and critic without requiring an additional disagreement ensemble. Exploration and planning guide training data collection, while evaluation uses the learned actor alone. Across Meta-World, Adroit, and DM-Control, SDI achieves a higher observed task-average score than every compared baseline in each suite and matches or exceeds the strongest baseline's performance on 15 of 18 tasks. Component ablations and robustness analyses further examine the roles of exploration, planning exposure, and action acceptance. Paired evaluations from identical simulator states show that the gate improves both recovery from actor failures and preservation of successful behavior relative to ungated planning, while substantially reducing return degradation. Together, these results support comparing imagined alternatives against a shared actor reference as a unified approach to exploring and planning with world models.

Figures & tables

Figure 1: Overview of SDI . Imagine: candidate action sequences (siblings, blue) and repeated rollouts of the actor’s fixed sequence (dashed gray) start at the same latent state

z_{t}

. Compare: fresh plan and actor rollouts provide the mean return gap

\hat{g}

and actor-return standard deviation

s

; dots represent candidate mean returns, whose spread is measured by

\mathrm{sib}

. Act/explore: the upper branch selects the planned action

\mu_{0}

or actor fallback

a^{\pi}

; dark cells indicate the executed action. The lower branch normalizes, clips, and weights the exploration bonus

b

to augment the first imagined reward

\tilde{r}

. The bonus anneals after the first training success. Exploration and gating use separate rollout settings; evaluation uses the actor alone.

Meta-World
Task	Plan2Explore		Curious Replay		TD-MPC2		DreamerV3		R2-Dreamer		SDI
	half	full	half	full	half	full	half	full	half	full	half	full
box-close	.00	.00	.67	.98	.81	.98	.00	.22	.58	.97	.90	.99
sweep-into	.00	.12	.77	.94	.76	.95	.00	.00	.90	.97	.79	.98
hammer	.28	.23	.00	.00	1.00	.99	.00	.00	1.00	1.00	.53	1.00
bin-picking	.00	.00	.00	.48	.22	.86	.08	.57	.48	.61	.01	.92

Table 1: Cross-domain benchmark performance under actor-only evaluation. Meta-World ( Yu et al., 2020a ) and Adroit ( Rajeswaran et al., 2018 ) report success; DM-Control ( Tassa et al., 2018 ) reports episode return. Compared baselines are Plan2Explore ( Sekar et al., 2020 ) , Curious Replay ( Kauvar et al., 2023 ) , TD-MPC2 ( Hansen et al., 2024 ) , DreamerV3 ( Hafner et al., 2025 ) , and R2-Dreamer ( Morihira et al., 2026 ) . Columns show half and full training budgets, and suite means weight tasks equally.

Figure 2: Cross-domain benchmark learning curves. Lines show average rolling success on Meta-World ( Yu et al., 2020a ) and Adroit ( Rajeswaran et al., 2018 ) and average episode return on DM-Control ( Tassa et al., 2018 ) . Black triangles mark the Table 1 scoring depths; curves for the latter suite may extend toward the nominal

500

k-step budget.

Ablation	Component	hand-insert		assembly		disassemble		mean	occupancy
		350k	700k	350k	700k	350k	700k	700k	mean
	SDI	.11	.49	.07	.75	.00	.35	.53	.97
Base agent	R2-Dreamer	.04	.13	.00	.00	.01	.12	.08	.00
Base agent	DreamerV3	.05	.16	.00	.00	.00	.00	.05	.37
Planning	Ungated planner	.20	.53	.00	.29	.03	.14	.32	.97
Planning	Half occupancy	.14	.47	.00	.00	.00	.01	.16	.47

Table 2: Ablation study on precision-manipulation tasks. Entries report mean evaluation success at

350

k and

700

k environment steps; the mean column gives the equally weighted average of the three

700

k-step task means. Occupancy is the fraction of acting steps with planner availability, not planned-action execution. Base-agent controls are R2-Dreamer ( Morihira et al., 2026 ) and DreamerV3 ( Hafner et al., 2025 ) .

Figure 3: Paired-state planning experiments. (a) Realized one-decision improvement over the actor by normalized-margin decile; the dashed line is the all-plan rate. (b) Proposal acceptance under four rules. (c) Recovery and success preservation during

50

-step continuations from restored states. (d) Checkpoint mean return difference between gated and ungated continuations. Shading in (a) shows the standard error; intervals in (c,d) are

95\%

World Model-Based Planning

World model-based planning uses models of how environments change to compare imagined actions and choose goal-directed behavior.

See all 261 papers

Jan 13, 2026

Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

Recent advances in world models have shown promise for modeling future dynamics of environmental states, enabling agents to reason and act without accessing real environments. Current methods mainly perform single-step or fixed-horizon rollouts, leaving their potential for complex task planning under-exploited. We propose Imagine-then-Plan (ITP), a unified framework for agent learning via lookahead imagination, where an agent's policy model interacts with the learned world model, yielding multi-step ``imagined'' trajectories. Since the imagination horizon may vary by tasks and stages, we introduce a novel adaptive lookahead mechanism by trading off the ultimate goal and task progress. The resulting imagined trajectories provide rich signals about future consequences, such as achieved progress and potential conflicts, which are fused with current observations, formulating a partially observable and imaginable Markov decision process to guide policy learning. We instantiate ITP with both training-free and reinforcement-trained variants. Extensive experiments across representative agent benchmarks demonstrate that ITP significantly outperforms competitive baselines. Further analyses validate that our adaptive lookahead largely enhances agents' reasoning capability, providing valuable insights into addressing broader, complex tasks. Our code and data will be publicly available at https://github.com/loyiv/ITP.

Youwei Liu, Jian Wang, Hanlin Wang +2

Central South University · Department of Computing, The Hong Kong Polytechnic University · College of Computer Science, Sichuan University

World Model-Based Planning Model-Based Planning Long-Horizon Agent Tasks LLM Agent Planning

Ask PDF

Oct 19, 2026

AnyGoal: On World Model Generalization Beyond Expert Rollouts

The goal of a world model is to capture the dynamics of an environment. This has the great appeal that it allows agents to plan and simulate a course of actions with the help of a world model. For these uses to be reliable, however, the learned dynamics must behave well beyond the specific state-action combinations observed in expert rollouts. Despite this, the relationship between data coverage, dynamics generalization, and downstream planning performance remains underexplored. We study this relationship through the lens of compositional generalization: can a world model recombine familiar state and action factors into valid dynamics? To answer this, we introduce AnyGoal, a benchmark of goal-shifted task variants built from widely used world-model environments, together with an open-source library. Each task is designed to be within reach of a world model that has learned the underlying dynamics: individual state factors and actions remain within their training support, but their combinations are underrepresented in expert trajectories. Across multiple environments and world model families, we find that current models consistently exhibit failures under these recombinations. By examining dynamics and representation failures across model families, interaction regimes, and evaluation metrics, we study where current world models generalize and where challenges remain for goal-shifted prediction and planning. Together, this framework, benchmark, and analysis provide a basis for evaluating progress toward world models that generalize to new planning tasks.

Heejeong Nam, Divyat Mahajan, Quentin Le Lidec +2

Brown University · Cold Spring Harbor Laboratory · Mila, Université de Montréal +1

World Models World Model-Based Planning World Model Evaluation

Ask PDF

Jun 1, 2026

IMWM: Intuition Models Complement World Models for Latent Planning

Planning with a learned latent world model is a promising route to control from raw pixels, but a strong world model alone is not enough. We show this experimentally: even with a perfect world model (operationalized by replacing the learned forward predictor with an idealized rollout of the true environment dynamics), a finite-budget sample-based planner still fails on some tasks, indicating that the bottleneck can lie in search rather than in world-model accuracy. Motivated by this gap, we propose IMWM (Intuition Model + World Model), which pairs the world model with an intuition model trained from demonstrations to recognize promising actions. The two models collaborate through three lightweight components: (i) Retrieval Initialization, which initializes the planner's action proposal from a retrieved demonstration; (ii) Hybrid Cost, which combines the intuition score with the world-model rollout cost; and (iii) a Reliability Gate, which adjusts how much the planner trusts intuition in each setting. Across four pixel-based goal-reaching tasks (Two-Room, Reacher, Push-T, and OGBench-Cube), IMWM has higher mean success than the world-model-only planner on all four, with the largest gains on Two-Room (99.2%, +11.5 percentage points) and OGBench-Cube (94.7%, +28.5 percentage points).

Baoqi Gao, Ruize Han, Miao Wang +1

Beihang University · Shenzhen University of Advanced Technology

World Model-Based Planning Goal-Conditioned Planning AI Control Model-Based Planning Latent World Models Latent-Space Planning

Ask PDF

Abstract

World Model-Based Planning

Explore similar work

Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

AnyGoal: On World Model Generalization Beyond Expert Rollouts

IMWM: Intuition Models Complement World Models for Latent Planning