LePlanner: An Iterative Amortized Controller For World Models
Authors: Saksham Bansal, Om Naphade, Chayan Aggarwal, Vrishin M
Organizations: Department of Mechanical Engineering, Indian Institute of Technology Roorkee · Department of Physics, Indian Institute of Technology Roorkee · Department of Mathematics, Indian Institute of Technology Roorkee · Department of Chemical Engineering, Indian Institute of Technology Roorkee
World models trained with joint-embedding predictive architectures learn compact, structured latent representations from physical interaction, yet planning in these latent spaces typically relies on one of two costly approaches. Search-based planners such as CEM, MPPI, and iCEM optimize action sequences through many predictor rollouts, achieving strong performance at the cost of high per-decision compute and latency. Policy-based methods amortize inference into a single forward pass but can degrade on contact-rich tasks where the demonstration distribution is multimodal. We propose LePlanner, an amortized iterative controller that learns to construct and refine latent action sequences through a frozen world-model predictor. LePlanner is trained with an arrival-and-hold objective that encourages the controller to reach the goal at the earliest feasible horizon and remain there. This addresses horizon-reset procrastination, a failure mode in which repeated receding-horizon replanning continually postpones goal arrival. An additional action-Gaussian loss keeps generated actions near the support of the offline dataset. Across navigation, contact-rich manipulation, and continuous-control environments, LePlanner matches or exceeds search-based planners while requiring an order of magnitude fewer predictor evaluations and 3-49x lower wall-clock time per decision. It achieves success rates of 98% on PushT, 100% on Reacher, 100% on TwoRooms, and 92% on the OGBench Cube task. These results show that much of the structure discovered through online search can be amortized into a lightweight iterative policy, enabling fast, horizon-aware, nonlinear physical control without online optimization.
Figures & tables
Figure 1: LePlanner architecture. The frozen JEPA predictor (blue) rolls out imagined latents xˉ1,xˉ2,… from context (E(si)+bi) , where E(si) is the latent embedding of state si and bi is the executed action. Y(k) denotes the action plan at refinement step k , produced by the Refinement Transformer and decoded into real-environment action states by the action head hθ .
Figure 2: The LePlanner architecture. From a 3-frame observation history and a goal image oG , the frozen encoder E produces context tokens Ct and goal latent zG . Plan tokens Yk are decoded by the action head hθ into an action plan b1:Hk , which the frozen predictor P rolls out into imagined latents z~1:Hk . The consequence encoder Fθ attends jointly over Yk , the rollout z~k , the rollout-to-goal residual z~k−zG , and the predicted distance-to-goal dk with contexts and goal to produce consequence tokens Rk . The joint refiner Gθ then applies full self-attention jointly over Yk and Rk , with the conditioning tokens Ct,zG to predict a plan correction ΔYk , yielding Yk+1 . This repeats for K refinements; after the final iteration the first action block b1 is executed, a new observation is received, and the controller replans.
Figure 3: Success across all compared methods. One panel per environment; bars are coloured by method family. LePlanner (green, outlined) is shown at both receding horizons.
Figure 4: Wall-clock per episode against matched CEM , log scale. LePlanner is 49× and 17× faster on TwoRooms, 36× and 7× faster on PushT, 18× and 4× faster on Reacher, and 9× and 3× faster on OGBench-Cube, at m=1 and m=5 respectively, while reaching higher success in every one of these comparisons.
No online
No action
Multi-step
No test-time
Frozen, task-
Fixed
Method
search
supervision
plan
deadline
agnostic WM
inference cost
CEM / MPPI / iCEM
✗
✓
✓
✓
✓
✗
Gradient optimisation
✗
✓
✓
✓
✓
✗
GCBC / GCIQL / GCIVL
✓
✗
✗
✓
✓
✓
GC-IDM
✓
✗
✗
✗
✓
✓
TD-MPC2
✓
✓
✓
✓
✗
✓
Table 1: What each planner family requires. ✓ requirement satisfied, ✗ not satisfied, ∼ partially. LePlanner is the only row with no hard failure. The ∼ marks that dataset actions bound LePlanner’s support region but are never regression targets.
Figure 5: Held-out imagined distance across the plan, split by the sample’s own goal offset q ; the open triangle marks each curve’s minimum. Left: under the original objective every curve bottoms out at block 5 , even the one whose goal is a single block away. Right: under arrival-and-hold the minimum tracks q , and for q=1 the curve rises afterwards, which is the signature of arriving and staying.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: The four evaluation environments. Top row: an observation ot drawn from the offline data. Bottom row: the goal oG used with it, which is the same episode 25 environment actions later — exactly the relabelling rule used at evaluation. Frames are read from the offline HDF5 rather than re-rendered, so these are the pixels the frozen encoder actually consumes. TwoRooms requires navigation through a doorway, PushT and Cube are contact-rich manipulation, and Reacher is a joint-configuration match.
Environment
Episodes
Frames
da
A
Offline dataset
PushT
18,685
2,336,736
2
10
pusht_expert
Reacher
10,000
2,010,000
2
10
reacher_expert
TwoRoom
10,000
920,809
2
10
tworoom
OGBench Cube
10,000
2,010,000
5
25
cube_single_expert
Appendix
Table 2: Offline datasets. One world-model transition is a block of f=5 environment actions in every environment, so the block width is A=5da and a plan of H=5 blocks is 25 environment actions. Frame counts are the rows in the released HDF5.
Table 3: The full stack. Only the controller receives gradient. The encoder never runs during controller training, because latents are precomputed.
Component
Definition
conditioning
N=3 context latents +1 goal latent, each projected 192→256 by one shared linear map, plus a learned positional embedding per context slot and a learned goal token. Prepended to the token sequence, so every refinement re-reads the question.
plan_query
H learned initial tokens, N(0,0.022) . Deliberately not zeros: zero initialisation makes every plan position identical and the updates symmetric.
transformer block
pre-norm: x+=MHA(LN(x)) then x+=MLP(LN(x)) , MLP ratio 4 , GELU, dropout 0.1 .
Fθ (consequence)
Linear(256+2⋅192+1→256) over [yj,z^j,z^j−zG,dj] , then 4 blocks, width 256 , 8 heads. It interprets the frozen rollout; it does not re-predict it.
Gθ (refinement)
Linear(2⋅256→256) , then 4 blocks, width 256 , 8 heads, full self-attention over all H plan tokens together with the conditioning tokens.
Δ head
Linear(256→256) , weights N(0,0.012) , bias 0 : refinements start near-identity while Fθ still receives gradient on the first pass.
Appendix
Table 4: Controller components. Fθ and Gθ are shared across all K refinements, so refinement depth buys compute and reasoning, never parameters.
Optimizer
AdamW, weight decay 10−4
Learning rate
3×10−4 (PushT, Reacher, Cube); 10−4 (TwoRoom)
Schedule
OneCycleLR, pct_start =0.05
Batch size
128 (PushT, Reacher, Cube); 32 (TwoRoom)
Steps
20,000
Gradient clipping
global norm 1.0
Goal curriculum
0:2, 0.25:3, 0.5:5 — max offset ≤2 , then ≤3 , then ≤5
Appendix
Table 5: Training configuration. Identical across environments except where noted.
Plan for a goal one block away
Profile d1…5
Old loss dH
New loss ( q=1 )
defers arrival
[0.090,0.053,0.035,0.023,0.014]
0.0140
0.1056
arrives and holds
[0.015,0.014,0.014,0.015,0.014]
0.0140
0.0221
Appendix
Table 6: The old objective is exactly indifferent. Two plans for a goal one block away, with their held-out distance profiles d1…5 . Scoring only dH cannot tell them apart; indexing arrival at q=1 can.
Figure 7: (a): The fitted closed loop under replan-every-block. Each replan removes a fraction 1−c of the remaining distance but adds a constant floor b , so the loop converges to D∗=b/(1−c) . Curves are evaluated from the measured (c,b) , not fitted by eye. The fixed point is a latent distance; its association with the 20 -pixel tolerance is inferred from the success rates that accompany it, not from a measured latent-to-pixel conversion. (b): refinement at eval time. Predicted cost improves through k=5 then worsens from k=6, while the plan keeps changing by ≈0.041 per iteration and never settles — the recurrent update is not contractive outside its trained depth.
k
terminal d5(k)
arrival dq(k)
gain Δ
mean plan change
0
0.05809
0.20163
—
0.00000
1
0.02641
0.16564
+0.03168
0.23859
2
0.01711
0.15428
+0.00930
0.13259
3
0.01486
0.15003
+0.00225
0.06835
Appendix
Table 7: Refinement converges within its trained depth. Held-out means across refinement iterations k . Terminal distance falls 3.9× from one pass to four, the gains shrink by an order of magnitude, and the plan moves less each time — matching the learned step sizes [0.550,0.451,0.312] . The pathology is visible here too: arrival stays ≈10× worse than terminal at every k , because the objective never asked for it.
Figure 8: The hold weight is not what fixes the failure. Left: success at both receding horizons. Every arrival-indexed setting (green) recovers; the fixed-terminal objective (red) collapses at rh=1 while remaining fine at rh=5 , which is the horizon-dependence itself. Right: the closed-loop fixed point D∗ falls monotonically with λh even where success does not, so the hold term is doing real work on drift that this task’s tolerance does not price.
Planner
rh=1
rh=5
Note
random actions
2%
—
floor
CEM ( 300 samples ×30 iterations)
34%
90%
≈45,000 predictor transitions per plan
LePlanner, fixed-terminal objective
52%
88%
horizon-dependent
LePlanner, arrival-and-hold
94%
88%
20 predictor transitions per plan
Appendix
Table 8: PushT scoreboard , 50 held-out episodes, identical start/goal pairs on every row. rh is the number of planned blocks executed before replanning. A healthy closed-loop planner should be roughly insensitive to rh ; the arrival-and-hold controller is ( 94 vs 88 , statistically tied), and the fixed-terminal controller is emphatically not.
Decoder
MSE
Mean-image MSE
Ratio
Held-out protocol
TwoRooms
5.97×10−6
1.93×10−3
324×
episode-split, 24,000 frames
PushT
6.78×10−4
4.82×10−3
7.1×
frame split of the cache
PushT
1.28×10−3
4.73×10−3
3.7×
256 fresh frames
Appendix
Table 9: Decoder fidelity. Both beat the mean-image baseline, but by very different margins. The TwoRooms decoder is a near-exact probe; the PushT decoder is a usable but visibly blurry one, so PushT panels should be read for object pose and gross geometry, not for fine texture.
Figure 9: Imagined versus actual rollout, TwoRooms, m=5 . Columns are plan blocks 1…5 ; the per-column figure is the predicted-versus-actual latent MSE. re-encoded tracks actual throughout, so the decoder is not the source of any disagreement. The world model is accurate for the first block (drift 0.101 ) and then diverges by roughly 16× (drift 1.85 onward): its imagined agent stalls against the dividing wall while the real agent passes through the doorway. The episode nevertheless succeeds, finishing at 15.5 pixels against a 16 -pixel tolerance — which is the case for replanning, not against it. Only block 1 is a complete block; blocks 2 – 5 are partial and padded, so their drift values are not directly comparable to block 1’s.
Figure 10: Imagined versus actual rollout, PushT. The same three rows on a contact-rich task. Here re-encoded is visibly blurrier than actual — that is the PushT decoder’s own error (Table 9 ), not the world model’s, and it is why this panel is read for pose rather than texture. Imagination degrades far more gently than on TwoRooms, with drift growing smoothly 0.009→0.029→0.040→0.063→0.074 across the five blocks rather than jumping an order of magnitude, and the imagined T-block pose stays close to the real one throughout.
Figure 11: What each refinement buys, TwoRooms. Top: the decoded imagined terminal state after refinement k . Bottom: the goal. The imagined endpoint migrates toward the goal as k increases while the predicted distance falls by roughly 7× over three refinements.
Figure 12: What each refinement buys, PushT. The initial plan k=0 is already close, and one refinement captures most of the available improvement ( 4× ), after which the plan is essentially settled — the behaviour a converging solver with decaying step sizes should show.
Figure 13: PushT: executed rollouts, four successful episodes. Each row is one episode: the state at each of ten replans, then the goal. The magenta outline is the target pose, and the final position and orientation error are given at the right. Note the first replans in each row, where the pusher travels around the block before it can move it. These episodes use sim-generated goals, not the held-out expert evaluation set that the success rates in Section 4 are computed on ; they illustrate behaviour and are not a success rate.
Figure 14: TwoRooms: executed rollouts, four successful held-out episodes at m=1 , controller seed 43 . Columns are evenly spaced replan points and the last column is the goal; the magenta ring is the 16 -pixel success tolerance, and the final physical distance is given at the right. Episodes differ in length, so rows with fewer replans are padded with blank space rather than repeated frames. These are the same held-out trials used throughout Appendix D , fixed before inspection.
Figure 15: Executed TwoRooms trajectory , m=1 . The agent path is drawn over the true floor plan with the goal tolerance ring; physical distance falls 44.8→15.0 pixels over the episode. Goal markers are drawn after rendering and never enter the encoder.
Modern vision-based world models can represent observations as compact yet expressive latent manifolds, but fast goal-oriented planning in these spaces remains challenging. This raises a central question: when does a learned representation simplify control, rather than merely enabling prediction? We study this question in a pretrained LeWorldModel, whose latent geometry is regularized for smoothness and uniformity. Our key insight is that, under such geometry, planning can be amortized into a latent inverse-dynamics mapping instead of requiring online search. We therefore replace iterative planning with a lightweight Goal-Conditioned Inverse Dynamics Model (GC-IDM) that maps the current latent state, goal latent state, and remaining horizon directly to the next action. Empirically, across four benchmark environments spanning navigation, contact-rich manipulation, and continuous control, our controller matches or exceeds CEM in seven of eight environment-protocol settings while reducing per-decision cost by 100-130x. A broader sweep over test-time planners (CEM, MPPI, iCEM, and gradient-based methods) shows that this result is not specific to a particular optimizer. These findings suggest that much of the structure recovered by test-time planning is already locally encoded in the latent representation. More broadly, our results indicate that sufficiently structured latent spaces can shift part of the planning burden from online optimization to learned inference. Our code is publicly available at https://github.com/hdnndh/Latent-Geometry-Beyond-Search-Amortizing-Planning-in-World-Models .
Hoang Nguyen, Xiaohao Xu, Xiaonan Huang
Department of Robotics, University of Michigan, Ann Arbor
Latent world models are inherently strong encoders that transform image pixel to latent embedding, yet existing world models still rely on online trajectory optimization for action planning: for every state-goal pair, an iterative optimizer is run from scratch to search for optimal action sequences, treating the world model as a black-box simulator. This approach pays the full iterative optimization cost anew at every replanning step and reuses no planning experience across queries. In this work, we ask whether planning itself can be amortized once a latent world model has been learned. We present LeFlow, which learns a reusable latent trajectory prior operating directly in the latent dynamics space from the world model. LeFlow recasts planning as conditional latent trajectory generation: a rectified-flow model imagines a future latent path between the current and goal embeddings, an inverse dynamics decoder turns latent transitions into action chunks, and the frozen world model verifies each candidate by autoregressive rollout. Across four major goal-conditioned pixel-control benchmarks, LeFlow replaces iterative action-space optimization with amortized latent planning and fixed-budget rollout selection, achieving consistent success-rate gains with an order-of-magnitude reduction in planning time. Our results argue that latent world models should support not only prediction but reusable planning priors. Our code is available at https://github.com/hsiangwei0903/LeFlow.
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