Knowing When to Stop Thinking: A Dissection of Expected Free Energy as a Test-Time Compute Controller for Latent World-Model Planning

Tong Zhang; Zhiyuan Shi; Yun Peng; Tao Xie

cs.ROOct 20, 2026

Knowing When to Stop Thinking: A Dissection of Expected Free Energy as a Test-Time Compute Controller for Latent World-Model Planning

Authors: Tong Zhang, Zhiyuan Shi, Yun Peng, Tao Xie

Organizations: Fudan University · Peking University

Abstract

A planner using a learned world model can improve its predicted outcome while choosing a worse action. We study when expected free energy (EFE) can help decide how much computation to spend on a plan. An EFE-inspired objective combines goal distance with disagreement among stochastic predictions. We test candidate scoring, stopping, and online predictor updates separately. Stopping when elite cost stalls reduces refinement on tasks whose success plateaus with depth, with mean success close to fixed-depth references. At the larger tested weights, rewarding disagreement lowers success with a frozen model. Penalizing it helps at some task and weight settings but harms others. When training data omit a region needed to reach the goal, online updates recover part of the resulting loss. The tested cost-stall rule can fail when further optimization exploits model errors, because predicted cost keeps falling after task success peaks. Cost progress can guide computation on a plateau, but protecting performance after a peak requires information about the consequences of the chosen actions.

Figures & tables

Panel of Figure 1 — Figure 1: Planning and measured effects of refinement. Top, the planning workflow refines candidate actions with a frozen world model until cost progress stalls or the budget is exhausted. Bottom, DINO-WM results from five seeds with 50 starts per seed. Success and model cost are plotted on separate axes. Each cost curve is divided by its first-iteration value within each seed before averaging. PushT success reaches a plateau. Wall success peaks at iteration 11 at $57.6\%$ and falls to $35.6\%$ at iteration 30 as model cost continues to decrease. All plotted curves come from recorded evaluations.

Figure 2: Planning and online adaptation. The upper loop encodes images, samples actions, predicts and scores futures, selects elites, and refits the proposal. The stopping rule either continues refinement or executes the selected action. The lower branch updates the predictor using collected transitions, replay, and a parameter anchor. Shuffled-target updates and a frozen predictor provide controls.

Quantity	Setting	Quantity	Setting
Model parameters	18M	Image resolution	$224\times 224$
Latent dimension	192	CEM samples / elites	300 / 30
Horizon / action block	5 / 5	Refinement cap	30
Stall threshold	$10^{-2}$	Patience / minimum depth	2 / 3
EMA coefficient	0.5	Episodes per seed	50
Initial controller seeds	3 to 6	Dose-response seeds per cell	100

Table 1: Primary LeWM configuration. The same cost-stall settings are used across the four tasks. Dose-response and transfer experiments have their own stated seed counts and budgets. The normalized-weight grid uses one seed per setting.

Environment	Best fixed	Primary	Best-elite
	success @ depth	success @ depth	success @ depth
TwoRoom	$90.0$ @ $30$	$88.0$ @ $22.0$	$89.3$ @ $9.3$
PushT	$88.4$ @ $30$	$87.2$ @ $28.3$	$87.3$ @ $15.3$
Cube	$67.3$ @ $10$	$63.3$ @ $28.1$	$63.3$ @ $15.8$
Reacher	$72.0$ @ $30$	$74.7$ @ $22.3$	$74.0$ @ $9.3$

Table 2: Controller success (%, mean across seeds) and mean depth. The primary rule saves 6 to 27% of iterations on TwoRoom, PushT, and Reacher relative to the cap. Each seed contains 50 episodes. Seed counts for fixed, primary, and best-elite are

(3,5,3)

on TwoRoom,

(5,5,3)

on PushT, and

(3,3,3)

on Cube and Reacher. Fixed depth is selected from the evaluated grid. The seed sets differ between methods.

Setting	Environment	$\beta$	Success change (pp)
Reward	All four	$+10$	$-3.5$ to $-54.7$
Reward	All four	$+50$	$-13.0$ to $-78.1$
Penalty	Reacher	$-20,-50$	$+2.38,+3.34$
Penalty	PushT	$-3$	$+1.42$
Penalty	Cube	$-20,-10$	$+1.44,+1.88$
Strong penalty	TwoRoom	$-20,-50$	$-4.7,-13.4$

Table 3: Success changes from the

\beta=0

baseline. Each setting uses 100 matched seeds and 50 episodes per seed with

K=4

stochastic rollouts. Differences are percentage points (pp). The reward rows give the range across tasks, and the penalty rows show selected positive and negative changes. Table 16 gives all weight comparisons.

Update	Hole e3	Hole e5	Matched	Full e2	Released
Seeds	5	5	5	5	3
Frozen	$37.2$	$44.0$	$37.2$	$70.8$	$75.3$
Online	$56.0$	$62.8$	$34.8$	$76.4$	$74.0$
Sham	$46.4$	$55.6$	$38.0$	$80.8$	n/a
Online $-$ sham (pp)	$+9.6$	$+7.2$	$-3.2$	$-4.4$	n/a
Deficit recovered	$37.6\%$	$48.0\%$	n/a	n/a	n/a

Table 4: Coverage intervention on a second independently trained sample. Success is the mean percentage over 50 episodes per seed. Checkpoints e2, e3, and e5 denote training epochs. The matched control is a random-drop e2 checkpoint matched to hole e3 on frozen success. Deficit recovery uses volume-matched e3/e5 references with success

87.2\%/83.2\%

. The update conditions in each column share a seed set.

Planner setting	Success	Mean depth
Fixed $K=1$ , curve peak	$0.576$	11
Fixed $K=1$ , full budget	$0.356$	30
Cost-only halt, $K=1$	$0.364$	29.04
Fixed $K=5$	$0.508$	30

Table 5: DINO-WM Wall with a cap of 30 iterations. Success is the mean over five seeds with 50 starts per seed. The peak is selected from the common-depth mean curve. The fixed-depth

K=5

setting uses more computation per iteration.

Family and task	Performance	Mean depth
TD-MPC2, 19 tasks	$98.9\%$ mean retained return	$3.8$
TD-MPC2, lowest retention	$96.1\%$ retained return	n/a
DINO-WM PushT, fixed	$0.852$ success	$30$
DINO-WM PushT, halt	$0.852$ success	$26.44$

Table 6: Cross-family stopping results. TD-MPC2 averages tasks equally and normalizes return by each task’s best fixed-depth mean on the same evaluation grid. DINO-WM uses five seeds and cap 30. Detailed task results and configurations appear in Tables 8 and 9 .

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Fixed-depth success and adaptive halting plotted against mean refinement depth. Figure 2 shows the planning procedure.

Panel of Figure 4 — Figure 4: Elite-cost trajectories and allocated depth. Left, example trajectories show rapid early descent and a slower-descending case that receives more iterations. Red markers indicate halting. Right, initial elite cost correlates with depth on PushT ( $r{\approx}0.3$ ). Initial cost is a model-based difficulty proxy. The correlation does not use an independently measured difficulty label.

Figure 5: Disagreement-weight sweep with 100 seeds per cell and 50 episodes per seed. (a) Mean success against $\beta$ on a symlog axis. Larger positive weights reduce success in every environment. Some negative weights give small gains on Reacher, PushT, and Cube. Strong penalties harm TwoRoom and PushT. (b) Mean paired differences from $\beta{=}0$ (Table 16 ). The gains depend on the environment and weight. Coverage is not independently varied.

Figure 6: DINO-WM depth curves with five seeds and a cap of 30. Wall’s mean success peaks at $0.576$ at iteration 11. It falls to $0.356$ at iteration 30, a $38.2\%$ relative decline. The peak is selected from the mean curve at a common depth across seeds. Cost-only halting with $\varepsilon=10^{-2}$ attains $0.364$ success at mean depth $29.04$ . The fixed-depth $K=5$ ensemble attains $0.508$ . The ensemble uses more stochastic rollouts at a fixed depth.

Figure 7: Checkpoint quality and uncertainty penalties. Fixed-depth success at $\beta{=}0$ and $\beta{=}{-}10$ is plotted against held-out one-step MSE on a log axis. Means are listed in Table 7 . TwoRoom’s curves remain close across the measured range. Reacher has positive contrasts at several checkpoints and a negative contrast at e8. Coverage is not independently varied. Stars mark released-model anchors with fewer seeds than the separate $n{=}100$ weight sweep.

Figure 8: Candidate preference and compute allocation. Positive $\beta$ favors terminal rollout spread and negative $\beta$ penalizes it. Highlighted bundles show candidate selection. Penalties can help or hurt task success. Cost-only halting ( $\beta=0$ ) controls refinement depth. Disagreement is an uncalibrated surrogate. Schematic.

Env	Ckpt	Probe MSE	$\beta{=}0$	$\beta{=}{-}10$	adaCost	adaEFE
TwoRoom	e1	$2.161$	$48.0$	$48.8$	$45.3$	$42.7$
	e2	$0.232$	$88.0$	$88.0$	$87.3$	$89.3$
	e3	$0.157$	$90.0$	$88.7$	$90.7$	$92.0$
	e5	$0.158$	$88.0$	$90.0$	$90.0$	$90.7$
	e8	$0.079$	$89.3$	$88.7$	$89.3$	$89.3$
	e12	$0.037$	$87.3$	$84.7$	$85.3$	$87.3$

Appendix

Table 7: Within-environment checkpoint comparison. Per-epoch LeWM checkpoints, 50 episodes per cell. Probe MSE,

\beta{=}0

\beta{=}{-}10

success, and adaptive halting. See Fig. 7 for visualization.

Task	$N^{*}$	Fixed return	Halt return	Relative return	Depth
acrobot-swingup	6	$661.2$	$661.8$	$100.1\%$	$3.1$
cheetah-run	9	$861.1$	$828.1$	$96.2\%$	$3.5$
cup-catch	9	$981.1$	$979.1$	$99.8\%$	$3.0$
dog-walk	9	$799.2$	$783.8$	$98.1\%$	$5.3$
finger-spin	6	$988.3$	$985.7$	$99.7\%$	$3.8$
finger-turn-hard	3	$972.7$	$967.0$	$99.4\%$	$3.0$

Appendix

Table 8: TD-MPC2 evaluation on nineteen DMControl tasks. Each task uses one released checkpoint and four seeds

\{2,3,4,5\}

, with 30 episodes per seed. Returns are means across seeds. The fixed reference maximizes mean return over

N\in\{3,6,9,12\}

on the evaluation seeds. Halting uses

\varepsilon=3\times 10^{-3}

. Its cap is 8 on dog-walk, humanoid-run, and humanoid-walk, and 6 on the remaining tasks. Relative return is

100\times\text{halt mean}/\text{fixed mean}

, so values above

100\%

indicate greater return than the fixed reference. Depth is averaged over seeds and summary rows weight tasks equally. Calibration evaluations are excluded. The fixed reference is selected retrospectively.

Configuration	PushT success @ depth	Wall success @ depth
Fixed $K=1$	$0.852$ @ $30$	$0.356$ @ $30$
Halt $\varepsilon=10^{-2}$	$0.852$ @ $26.44$	$0.364$ @ $29.04$
Halt $\varepsilon=3\times 10^{-3}$	$0.852$ @ $28.71$	$0.356$ @ $29.79$
Fixed $K=5$	$\cdot$	$0.508$ @ $30$

Appendix

Table 9: DINO-WM evaluation with five seeds (1 to 5), 50 starts per seed, and cap 30. Success is the mean across seeds. Depth counts the iterations used by each trajectory. The

K=5

configuration uses fixed-depth ensemble scoring.

	Success and fixed-depth gap			Coverage
Env.	Best $d$	$\Delta\mathrm{EFE}$	$\Delta$	$\mathcal{R}^{\mathrm{post}}_{\mathrm{model}}$
TwoRoom	$90.0$ @ $30$	$88.0$	$+2.0$	$\cdot$
PushT	$88.4$ @ $30$	$87.2$	$+1.2$	$\cdot$
Reacher	$72.0$ @ $30$	$74.7$	$-2.7$	$\cdot$
Cube	$67.3$ @ $10$	$63.3$	$+4.0$	$\cdot$
hole (e3)	$50.0\to 31.2$ (TTT $+18.8$ , $38\%$ )			$31.2$

Appendix

Table 10: Stopping and online adaptation. Above, success relative to fixed-depth and coverage references. Below, the TwoRoom hole-e3 factorial uses seeds

\{42,7,123,5,99\}

and 50 episodes per seed. Cells report mean success. The mean interaction is

+0.4

pp, with within-seed values

\{-6,-4,+4,0,+8\}

pp. The factorial and Table 4 are separate evaluations of the same labeled condition.

	TwoRoom	Reacher	PushT	Cube
Seeds with full traces	$1,2,3,4$	$1,2$	$3,4$	$1,2$
Calls with full traces	$241$	$156$	$124$	$142$
Calls without full traces	$62$	$81$	$196$	$65$
Mean $r_{\mathrm{ps}}$	$6.395$	$2.723$	$0.075$	$0.042$
Calls with positive denominator	$239$	$151$	$124$	$141$
Mean $\widetilde{r}_{\mathrm{ps}}$	$0.178$	$0.247$	$0.007$	$0.015$

Appendix

Table 11: Primary-controller elite-summary regret. Only complete trajectories in the listed seed subsets contribute to the cost averages. All calls are counted regardless of episode outcome. The normalized row uses iteration 3 as its reference and excludes eight zero denominators, seven with positive raw regret. The trace subset differs from the full-seed episode-success comparison.

TwoRoom, sample 1	Frozen	Online	Change
Coverage hole	$54.7$	$69.3$	$+14.7$
Volume control	$66.0$	$70.7$	$+4.7$
Full data	$74.0$	$75.3$	$+1.3$
Frozen penalty contrast	Sample 1		Sample 2
Coverage hole	$+20.0$		$-4.0$
Reacher	Online change		Interaction

Appendix

Table 12: Additional online-adaptation results. Success and changes are percentages and percentage points, respectively. Paired changes are reported separately from rounded success means. Sample 1 uses three paired seeds per TwoRoom condition. Its hole-model gain is positive in all three. Sample 1 Reacher uses five paired seeds, with four negative changes at e3. The separate sample 2 volume-matched references have frozen success

87.2

at e3 and

83.2

at e5. On collected Reacher transitions, prediction error decreases from approximately

1.1

0.63

under adaptation, and anchor error stays flat.

Training sample	Checkpoint	Probe MSE	Penalty change
1	e3	$0.194$	$+6.8$
2	e5	$0.122$	$+8.0$
2	e8	$0.035$	$+2.0$
2	e2, forward test	$0.131$	$+3.2$

Appendix

Table 13: Reacher probe error and frozen-penalty changes in success. Differences are percentage points. The sample 2 forward test at e2 has two positive seeds out of five and misses its prespecified criterion.

Halting signal	TwoRoom	PushT	Cube	Reacher
$\Delta\mathrm{EFE}$ EMA-mean (primary)	$88.0$ @ $22.0$	$87.2$ @ $28.3$	$63.3$ @ $28.1$	$74.7$ @ $22.3$
no-EMA (v1)	$87.3$ @ $16.4$	$\cdot$	$\cdot$	$\cdot$
epistemic-aware (adaEFE)	$90.0$ @ $24.9$	$96.0$ @ $29.8$	$72.0$ @ $29.2$	$78.5$ @ $26.7$
best-elite cost stall	$89.3$ @ $9.3$	$87.3$ @ $15.3$	$63.3$ @ $15.8$	$74.0$ @ $9.3$
elite-variance stall	$88.0$ @ $8.5$	$77.3$ @ $20.9$	$66.0$ @ $10.3$	$76.7$ @ $19.0$
mean-change stall	$88.7$ @ $22.5$	$89.3$ @ $17.9$	$65.3$ @ $24.7$	$76.0$ @ $26.2$

Appendix

Table 14: Stopping-rule comparison. Success is the mean percentage across seeds. The second line in each adaptive cell gives mean depth. Each seed uses 50 episodes and cap 30. Seed sets and thresholds are specified below. Mean differences are descriptive where seed sets differ. Missing entries indicate unevaluated configurations.

Figure 9: Success and refinement depth for the stopping rules. Each signal is evaluated across its threshold grid. The gray line shows fixed-depth success and the black star marks the primary controller. Cost-based rules follow the plateau with different depth allocations. Geometric rules are sensitive to thresholds. No evaluated rule reliably recovers Cube’s earlier success peak.

Rule	Cube	TwoRoom	PushT	Reacher
fixed-10 (Cube peak)	$71.3$	$\cdot$	$\cdot$	$\cdot$
fixed-30	$61.0$	$\cdot$	$\cdot$	$\cdot$
primary, unguarded	$67.0$ @ $28.3$	$\cdot$	$\cdot$	$\cdot$
guard $\gamma{=}1.15$	$68.7$ @ $18.1$	$\cdot$	$\cdot$	$\cdot$
guard $\gamma{=}1.3$	$64.0$ @ $20.5$	$82.0$ @ $12.4$	$69.3$ @ $12.1$	$78.7$ @ $21.4$
guard $\gamma{=}1.5$	$66.7$ @ $22.3$	$\cdot$	$\cdot$	$\cdot$

Appendix

Table 15: Disagreement guards. Mean success is followed by mean depth. Each cell uses three seeds and 50 episodes per seed. The guard stops when disagreement exceeds

\gamma

times its running minimum. It supplements the primary controller. The plateau tasks use

\gamma=1.3

. At these seeds, the unguarded adaptive reference reaches

86

on TwoRoom and

92

98

on PushT. Fixed-30 reaches

76

on Reacher. Missing entries indicate unevaluated configurations.

Env	$\beta$	Success (%)	$\Delta$ (pp)
TwoRoom	-50	$71.06$	$-13.44$
TwoRoom	-20	$79.82$	$-4.68$
TwoRoom	-10	$82.36$	$-2.14$
TwoRoom	-3	$83.12$	$-1.38$
TwoRoom	0	$84.50$	$\cdot$
TwoRoom	+3	$76.42$	$-8.08$

Appendix

Table 16: Disagreement-weight comparison with 100 paired seeds and 50 episodes per run. Success is the mean percentage across seeds.

\Delta

is the mean difference from the same-seed

\beta=0

reference. Every configuration uses 30 CEM iterations and

K=4

stochastic rollouts.

Figure 10: Disagreement-weight results with 100 seeds per cell. (a) Mean success. (b) Mean paired differences from $\beta=0$ . (c) Ratio of mean success at $\beta=0$ to $+50$ , reaching $10.54\times$ on PushT. Each seed uses 50 episodes, giving 160,000 episode evaluations across the weight grid. Reacher’s $\beta=-20,-50$ differences are $+2.38,+3.34$ pp, compared with $+4.4$ pp at each weight in the earlier 20-seed sample. In a separate set of repeated evaluations, Reacher success was $90\%$ initially and $80$ to $82\%$ in five repetitions.

Figure 11: Scale-matched dose response by environment. Seven normalized weights and 28 values use a common scale from $0$ to $100\%$ . Baseline annotations give mean success. Nonzero-weight cells use a single seed. Disagreement magnitude is matched to the current candidate batch. The intermediate positive weight reduces success by $20$ to $52\%$ relative on three environments. The intermediate negative weight improves Reacher success by approximately $11\%$ relative.

	$K{=}2$	$K{=}4$	$K{=}8$	$p{=}0.05$	$p{=}0.1$	$p{=}0.2$
Success	$80$	$80$	$82$	$76$	$80$	$64$

Appendix

Table 17: Proxy sensitivity on Reacher at

\hat{\beta}=-0.3

. The ensemble-size sweep fixes dropout at

p=0.1

. The dropout sweep fixes

K=4

. Each cell uses seed 42. The separate multi-seed

\beta=0

baseline is

72.0

$\hat{\beta}$ (spread-normalized)	$-0.3$	$-1$	$+0.3$	$+1$
Return, % of $\beta{=}0$	$99.2\%$	$96.5\%$	$32.9\%$	$5.9\%$
$n$ tasks complete	$2$	$10$	$11$	$10$

Appendix

Table 18: Spread-normalized weights on TD-MPC2. Raw

\beta

equals

\hat{\beta}

times per-task candidate-value spread divided by disagreement scale. Depth is fixed at 6. Return is a percentage of the

\beta=0

baseline. Each column covers

n

of 19 tasks, with different task subsets. Value magnitude is approximately

90\times

candidate spread. The spread contracts approximately

70\times

over six MPPI iterations. Dog-walk reaches

117.2\%

\hat{\beta}=-1

in one 10-episode evaluation.

Expected Free Energy

See all 10 papers

The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use
Aug 13, 2026Joyjeet SinghWorld Model-Based Planning Model-Based Planning
Independent researcher
Expected Free Energy-based Planning as Variational Inference
Jun 9, 2026Wouter W. L. Nuijten, Thijs van de Laar, Bert de VriesFree Energy Calculation Active Inference
Eindhoven University of Technology, Eindhoven, the Netherlands · Lazy Dynamics B.V., the Netherlands
Not All Errors Matter: Decision-Relevant Prediction Error Predicts Planning Quality
Sep 26, 2026Linhao Wang, Yiyan Fan, Dongjin HuangWorld Model-Based Planning Model-Based Planning
Shanghai Film Academy, Shanghai University

Aug 13, 2026J/K move · Enter open · S save

The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use

Joyjeet Singh

Independent researcher

World Model-Based Planning Model-Based Planning Latent World Models World Model Evaluation Latent-Space Planning

Open paper AskPDF arXiv

Latent world models are judged by how well they predict, so when planning fails at long horizons the natural reading is that the predictor degrades. On a reproduction of LeWorldModel on TwoRoom we show the binding constraint is the planner's objective instead. The predictor is not the limit: its imagined state seventy-five environment steps ahead is still only 0.189 as wrong as assuming the world froze, while the planner never imagines beyond twenty-five. The objective is. Cross-entropy-method planning minimises squared latent distance, which tracks true distance at r = 0.426, saturates by about eighty arena units and decreases beyond a hundred and twenty, so moving away from the goal can lower the cost. The information is present throughout: a ridge probe recovers position from the frozen embedding at R^2 0.9922. The pathology is the method's, not one reimplementation's. It is present in the authors' released weights, and across four checkpoints long-horizon success rank-orders exactly with metric quality and inversely with prediction accuracy. Replacing only the objective, with nothing retrained and no GPU, lifts goals reached at offset 100 from 26.0% to 98.0%, equals the 98.0% at offset 25, and reaches 92.0% under a third of the budget: planning stops depending on the horizon. The best cost is not the most accurate. A head learned from frame separation alone predicts spatial distance worse than a position probe (r = 0.819 against 0.9897) yet plans better, charging 24% more to cross the environment's dividing wall where squared latent distance charges 4% less. It has learned reachability, not proximity.

Abstract

Expected Free Energy

Explore similar work

The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use

Expected Free Energy-based Planning as Variational Inference

Not All Errors Matter: Decision-Relevant Prediction Error Predicts Planning Quality