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
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×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
β
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 β=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.
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≈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 β 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 β=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 ε=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 β=0 and β=−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 β favors terminal rollout spread and negative β penalizes it. Highlighted bundles show candidate selection. Penalties can help or hurt task success. Cost-only halting ( β=0 ) controls refinement depth. Disagreement is an uncalibrated surrogate. Schematic.
Env
Ckpt
Probe MSE
β=0
β=−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, β=0 / β=−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∈{3,6,9,12} on the evaluation seeds. Halting uses ε=3×10−3 . Its cap is 8 on dog-walk, humanoid-run, and humanoid-walk, and 6 on the remaining tasks. Relative return is 100×halt mean/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 ε=10−2
0.852 @ 26.44
0.364 @ 29.04
Halt ε=3×10−3
0.852 @ 28.71
0.356 @ 29.79
Fixed K=5
⋅
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
ΔEFE
Δ
Rmodelpost
TwoRoom
90.0 @ 30
88.0
+2.0
⋅
PushT
88.4 @ 30
87.2
+1.2
⋅
Reacher
72.0 @ 30
74.7
−2.7
⋅
Cube
67.3 @ 10
63.3
+4.0
⋅
hole (e3)
50.0→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 rps
6.395
2.723
0.075
0.042
Calls with positive denominator
239
151
124
141
Mean rps
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 to 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
Δ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
⋅
⋅
⋅
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
⋅
⋅
⋅
fixed-30
61.0
⋅
⋅
⋅
primary, unguarded
67.0 @ 28.3
⋅
⋅
⋅
guard γ=1.15
68.7 @ 18.1
⋅
⋅
⋅
guard γ=1.3
64.0 @ 20.5
82.0 @ 12.4
69.3 @ 12.1
78.7 @ 21.4
guard γ=1.5
66.7 @ 22.3
⋅
⋅
⋅
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 γ times its running minimum. It supplements the primary controller. The plateau tasks use γ=1.3 . At these seeds, the unguarded adaptive reference reaches 86 on TwoRoom and 92 to 98 on PushT. Fixed-30 reaches 76 on Reacher. Missing entries indicate unevaluated configurations.
Env
β
Success (%)
Δ (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
⋅
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. Δ is the mean difference from the same-seed β=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 β=0 . (c) Ratio of mean success at β=0 to +50 , reaching 10.54× on PushT. Each seed uses 50 episodes, giving 160,000 episode evaluations across the weight grid. Reacher’s β=−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 β^=−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 β=0 baseline is 72.0 .
β^ (spread-normalized)
−0.3
−1
+0.3
+1
Return, % of β=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 β equals β^ times per-task candidate-value spread divided by disagreement scale. Depth is fixed at 6. Return is a percentage of the β=0 baseline. Each column covers n of 19 tasks, with different task subsets. Value magnitude is approximately 90× candidate spread. The spread contracts approximately 70× over six MPPI iterations. Dog-walk reaches 117.2% at β^=−1 in one 10-episode evaluation.
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.
Planning under uncertainty requires agents to balance goal achievement with information gathering. Active inference addresses this through the Expected Free Energy (EFE), a cost function that unifies instrumental and epistemic objectives. However, existing EFE-based methods typically employ specialized optimization procedures that are difficult to extend or analyze. In this paper, we show that EFE-based planning can be formulated as Variational Free Energy minimization on a generative model augmented with epistemic priors. Our main result demonstrates that minimizing a Variational Free Energy functional with appropriately chosen priors yields a decomposition into expected plan costs (the EFE) plus a complexity term. This formulation reinforces theoretical consistency with the Free Energy Principle by casting planning as the same inferential process that governs perception and learning. We validate our approach on three environments of increasing complexity: a deterministic T-maze, a stochastic Reactivity Maze, and a partially observable MiniGrid DoorKey-8x8 environment. The experiments demonstrate that the epistemic priors induce information-seeking behavior, that the variational formulation yields policy-based inference outperforming plan-based methods under stochastic transitions, and that temporal factorization enables scalability to environments where existing tabular active inference methods cannot operate.
Wouter W. L. Nuijten, Thijs van de Laar, Bert de Vries
Eindhoven University of Technology, Eindhoven, the Netherlands · Lazy Dynamics B.V., the Netherlands
World models are typically trained and evaluated by prediction error, assuming that more accurate predictions lead to better decisions. We show that this assumption can fail because models with similar total error can differ substantially in planning performance when their errors occur on different state dimensions. We introduce Decision-Relevant Prediction Error (DRPE), which measures prediction error on the state dimensions that affect decisions. We also develop an iso-error evaluation protocol that varies error allocation while keeping total error fixed. In a factored gridworld with known state relevance and a standardized planner, we evaluate 55 controlled and learned models across different error levels and allocations. Total prediction error is weakly related to planning success (Spearman ρ=−0.25), whereas DRPE is strongly predictive (ρ=−0.84; −0.98 within the controlled family). Models with only a 1% difference in total error can differ by 60 percentage points in planning success (97% vs 37%). The relevant error also depends on the task, with model rankings reversing across tasks at the same total error. Deeper imagination further amplifies decision-relevant errors, while learned models exhibit systematic bias on rare but decision-critical events. We formalize sufficient conditions under which DRPE correctly ranks models and total prediction error cannot.