PaperlayerView exact revision on arXiv

Source-linked AI summary

LePlanner: An Iterative Amortized Controller For World Models

Saksham Bansal, Om Naphade, Chayan Aggarwal, Vrishin M

arXiv:2609.13845v1cs.ROcs.AI

TL;DR

Planning in latent world models often requires costly online search, while amortized policies can struggle on contact-rich tasks; fixed terminal deadlines can also postpone goal arrival under repeated replanning. LePlanner iteratively refines plans through a frozen predictor using arrival-and-hold training and an offline action-support constraint. It outperforms the matched CEM baseline while using substantially fewer predictor evaluations, though current experiments cover only four simulated environments.

  • Problem

    Search-based latent planners incur high per-decision cost, cloning-based policies can degrade on contact-rich tasks, and repeated replanning can reset a terminal deadline and defer goal arrival.

    Source: §1 · pp. 1–2

  • Method

    LePlanner iteratively refines action plans through a frozen world model, training with an arrival-and-hold objective and an action-support constraint derived from offline data.

    Source: §1 · p. 2, §3.3 · p. 5, §3.5 · p. 5

  • Results

    LePlanner outperforms matched CEM across the measured environments while requiring substantially fewer predictor evaluations.

    Source: §4.1 · p. 6, §4.2 · p. 7, §5 · p. 8

  • Takeaways & Limitations

    A fixed number of learned refinements can preserve multi-step world-model reasoning while replacing costly online trajectory search.

    Source: §5 · p. 8

    LePlanner depends on frozen-world-model accuracy and offline-dataset coverage, and current experiments cover only four simulated environments.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Latent world models make planning in compact representations possible, but controllers face a trade-off between costly online search and policies that can struggle on contact-rich tasks. LePlanner combines iterative plan refinement through a frozen predictor with an arrival-and-hold objective to address both inference cost and delayed goal arrival.

1 Introduction

  • Search-based planners are costly and high-latency, while cloning-based policies can degrade on contact-rich and nonlinear tasks with poorly conditioned demonstrations.Model-based policy methods also couple the controller to a task-specific dynamics model and can struggle with complex long-horizon structure.Source: §1 · pp. 1–2
  • Repeated replanning can reset a fixed terminal deadline, causing a controller to approach the goal while postponing actual arrival.The paper calls this failure mode horizon-reset procrastination.Source: §1 · p. 2
  • LePlanner iteratively refines action plans through a frozen LeWM predictor and trains with an arrival-and-hold objective to encourage timely arrival and continued proximity to the goal.The sampled goal offset defines the training loss but is not provided to the controller at inference.Source: §1 · p. 2

2 Related Work

Prior latent-world planning approaches use online search, learned policies, or policies trained inside model imagination. LePlanner combines iterative refinement with a frozen world model and an explicit arrival-and-hold objective.

2 Related Work

  • LePlanner keeps the learned representation and dynamics fixed, focusing instead on how the controller objective interacts with receding-horizon execution.This makes the work complementary to methods that reduce planning-model mismatch or regularize latent trajectory geometry.Source: §2 · p. 3
  • Search planners optimize action sequences through learned dynamics, while amortized goal-conditioned policies map state, goal, and horizon directly to actions but can degrade on contact-rich or multimodal tasks.Dreamer and TDMPC2 instead train policies within imagined rollouts.Source: §2 · p. 3
  • LePlanner combines multi-horizon reasoning through learned dynamics with a fixed number of learned refinements at deployment and an arrival-and-hold objective.Unlike model-based policy methods described here, it uses a frozen, task-agnostic world model.Source: §2 · p. 3
  • Unlike online CEM and expert-action imitation, LePlanner learns iterative corrections while using a pretrained, frozen world model.This separation allows the controller objective to change without changing the visual representation or dynamics.Source: §2 · p. 3

3 Method: LePlanner

LePlanner learns goal-conditioned action plans by repeatedly refining them through a frozen latent world model. Its training combines arrival-and-hold supervision, deep supervision across refinements, and an offline action-support constraint; deployment uses fixed learned refinements with receding-horizon execution.

3.1 Problem setting

  • The task is offline, visual, goal-conditioned control: produce action plans from observation history and a goal image without imitating dataset actions.A frozen encoder and predictor represent context, goal, and candidate-plan rollouts in latent space.Source: §3 · p. 3, §3.1 · p. 3

3.2 LePlanner

  • LePlanner initializes plan tokens, decodes them into bounded action blocks, rolls each plan through the frozen predictor, and iteratively revises the plan using predicted consequences.The refinement transformer jointly attends across plan positions; its operators are shared across iterations.Source: §3.1 · p. 4, §3.2 · pp. 4–5

3.3 Arrival-and-hold objective

  • The arrival-and-hold objective evaluates the predicted trajectory at the hindsight goal’s sampled offset and penalizes moving away afterward.The offset is used only to index the training loss, not supplied as a controller input at inference.Source: §3.3 · p. 5

3.4 Deep supervision across refinements

  • Deep supervision trains all K + 1 plan predictions, giving every intermediate plan a direct gradient signal and larger weights to later refinements.This trains the initial prediction to be useful while also training each shared refinement step.Source: §3.4 · p. 5

3.5 Constraining plans to offline support

  • A context-conditioned Gaussian mixture estimates offline action support, and a one-sided penalty constrains refined plans without using dataset actions as controller targets.The penalty is zero for action blocks inside the selected support region; only LePlanner parameters are optimized.Source: §3.5 · p. 5

3.6 Receding-horizon execution

  • At deployment, LePlanner performs K learned refinements and returns a plan whose prefix is executed before observing again and replanning.The inference loop retains the world model but replaces population-based search with a fixed number of amortized refinements.Source: §3.6 · p. 6

4 Results

Across four benchmark environments and several planner families, LePlanner combines strong goal-reaching success with lower matched-CEM decision cost. Its family comparison highlights distinct data and timing requirements, while held-out distance profiles show arrival-and-hold aligning predicted arrival with each goal’s sampled offset.

4.1 Goal-reaching success

  • LePlanner reached 100% on TwoRooms, 94% on PushT, 100% on Reacher, and 92% on OGBench-Cube, improving success over matched CEM in every measured environment.Its success stayed close to flat across replanning intervals, while CEM on PushT shifted from 34 to 90 between m=1 and m=5.Source: §4.1 · p. 6

4.2 The cost of a decision

  • LePlanner was 49×/17× faster on TwoRooms, 36×/7× on PushT, 18×/4× on Reacher, and 9×/3× on OGBench-Cube at m=1/m=5, with higher success in each comparison.It uses 20 latent transitions per plan versus CEM’s 45,000, a factor of 2250.Source: Figure 4 · p. 7, §4.2 · p. 7

4.3 What each family requires

  • LePlanner uses offline actions only to bound support through a frozen density, not as supervision targets; unlike LePlanner, GC-IDM requires the remaining test-time budget.Table 1 identifies LePlanner as the only planner-family row without a hard failure; its dataset-action support constraint is partial rather than a regression target.Source: §4.3 · p. 7, Table 1 · p. 7

4.4 The mechanism: arrival at the data’s own deadline

  • A fixed-terminal objective can postpone arrival indefinitely as replanning resets its deadline; arrival-and-hold instead puts each predicted-distance minimum at the sampled goal offset q.For q=1, the distance rises after the minimum, indicating arrival followed by holding position.Source: §4.4 · p. 8, Figure 5 · p. 8

5 Discussion

LePlanner combines iterative planning through a frozen latent world model with an arrival-and-hold objective, rather than cloning dataset actions. Its scope remains bounded by world-model accuracy, offline-data coverage, and evaluation on four simulated environments.

Discussion

  • LePlanner plans through a frozen latent world model without costly online optimization, outperforming matched CEM while requiring substantially fewer predictor evaluations.Its arrival-and-hold objective encourages timely goal arrival and continued proximity, and stable performance across replanning intervals suggests time-consistent behavior.Source: §5 · p. 8
  • Unlike behavior cloning, LePlanner minimizes goal-reaching losses through the frozen world model rather than regressing onto demonstrated actions.Offline trajectories provide contexts and hindsight goals, while dataset actions estimate the action-support constraint rather than serving as targets.Source: §5 · p. 8
  • LePlanner depends on frozen-world-model accuracy and offline-data coverage, and the current evaluation is limited to four simulated environments.Prediction errors may accumulate over longer horizons, while limited coverage can restrict valid-plan discovery.Source: §5 · p. 8

A Experimental Details

LePlanner is evaluated on four visual-control tasks using offline datasets and frozen world models. The appendix specifies the evaluation protocol and the controller’s reusable architecture, including its behavior-density regularizer.

A Experimental Details

  • Four visual-control tasks use released offline datasets, cached 192-dimensional latents, and episode-level 95/5 training and validation splits.The tasks are PushT, Reacher, TwoRoom, and OGBench Cube; evaluation uses 50 seeded start-goal trials per environment and shared trial manifests for paired comparisons.Source: A Experimental Details · pp. 12–13
  • The reported PushT result is the best pair among eleven task-selection draws, not an unbiased estimate.No draw met the prespecified threshold; across draws, success ranged from 88–100% at m=1 and 78–96% at m=5.Source: A Experimental Details · p. 13

B.1 What is frozen and what is learned

  • The encoder and predictor stay frozen, while refinement modules Fθ and Gθ share parameters across iterations.The planning space is the 192-dimensional post-projector embedding, and precomputed latents keep the encoder out of controller training.Source: §B.1 · p. 13, Table 3 · p. 14, Table 4 · p. 14

B.2 Controller

  • The controller reasons over one hidden plan token per action block, decodes bounded actions for prediction, and aligns the current frame with the first plan block.The action bounds are exact in normalized units through tanh.Source: §B.2 · p. 14

B.3 Behaviour density

  • A conditional Gaussian mixture models dataset action blocks, and the support penalty is zero within the selected 95th-percentile region.For PushT, c95 is 1.5306; the penalty is evaluated across refinements and plan steps.Source: §B.3 · p. 14

B.4 Training

The training procedure iteratively refines action plans through a frozen predictor, combining arrival-and-hold supervision with a small support penalty. The controller solves for goal achievement rather than copying dataset actions.

B.4 Training

  • Each training step rolls bounded action blocks through the frozen predictor and refines the plan, minimizing a weighted arrival-and-hold loss plus 0.01 times the support hinge.The predictor remains differentiable so gradients can update the plan.Source: §B.4 · p. 15
  • Removing the support penalty left success statistically unchanged (50.0 vs 52.0, p = 1.0) but raised off-support blocks from 0.187 to 0.652.The support term contributes only 0.000208 to a representative total loss of about 0.0201.Source: §B.4 · p. 15
  • On 256 held-out samples, LePlanner’s terminal distance was 0.0211 versus 0.3930 for replaying the expert block, winning on 99.6% of samples.The controller’s first predicted block differed substantially from the dataset block, with mean cosine similarity 0.552.Source: §B.4 · p. 15

C The Arrival-and-Hold Objective: Diagnosis and Ablation

Under the fixed-terminal objective, frequent replanning underperformed full-plan execution, and measurements ruled out several suspected explanations. The diagnosis was horizon-reset procrastination: scoring only the fixed endpoint failed to reward prompt arrival, leaving the closed loop at an unsuccessful fixed point.

C The Arrival-and-Hold Objective: Diagnosis and Ablation

C.1 The symptom

  • With the original fixed-terminal objective, executing the full five-block plan reached 88% success versus 52% when replanning after every block.The discrepancy persisted across several seed checks, interaction budgets, and an independent ground-truth-space harness.Source: §C.1 · p. 15

C.2 Root cause: horizon-reset procrastination

  • For a goal one block away, the old loss scored two distinct distance profiles identically, while scoring arrival at q=1 preferred prompt arrival by 4.8×.Scoring only dH cannot distinguish the profiles; indexing arrival at q=1 can.Source: §C.2 · p. 15, Table 6 · p. 16

C.3 The closed loop

  • The original objective converged to a fixed point outside the success radius, so quadrupling the step budget did not improve the outcome.Under the fix, c fell from 0.579 to 0.392 and the floor b fell from 0.041 to 0.024.Source: §C.3 · p. 16

C.4 How much depth the refinement loop actually buys

Shared refinement modules improve plans within the trained loop depth, but additional test-time iterations do not continue that improvement. The recurrent update is not contractive beyond its trained depth.

C.4 How much depth the refinement loop actually buys

  • Refinement reuses the same Fθ and Gθ at every pass, so adding iterations increases computation without adding parameters.Source: §C.4 · p. 16
  • Terminal distance fell 3.9× from one pass to four, but predicted cost worsened after k=5 and the plan kept changing by ≈0.041 per iteration.The update was not contractive beyond the trained depth, so K=3 is a trained operating point rather than a free evaluation-time knob.Source: §C.4 · p. 16, Figure 7 · p. 16, Table 7 · p. 17
  • Within the measured refinement passes, plan changes shrink as learned step sizes decrease from 0.550 to 0.451 to 0.312.The table also reports that arrival remains ≈10× worse than terminal distance at every pass.Source: Table 7 · p. 17

C.5 Which part of the change did the work

Arrival indexing fixes the horizon-dependent failure, while the hold term reduces post-arrival drift without changing success on this task. The comparison also covers controller behavior across replanning horizons and the scope of the diagnostic decoders.

C.5 Which part of the change did the work

  • Re-indexing arrival by q fixes the pathology; the hold term reduces post-arrival drift, which remains within the 20-pixel tolerance.The final-distance-to-minimum ratio falls from 1.9–5.1× to 1.0–3.1×, without converting drift into failed episodes on this task.Source: §C.5 · p. 17, Figure 8 · p. 17
  • The support penalty leaves success unchanged but reduces off-support plans, while removing path loss worsens the pathology.Removing the support penalty changes off-support blocks from 0.187 to 0.652 with success 50.0 vs 52.0 (p = 1.0); pure dH yields 18.0% success at rh=1.Source: §C.5 · p. 17

C.6 Summary of the comparison

  • 94 vs 88: the arrival-and-hold controller is statistically tied across replanning horizons, unlike the fixed-terminal controller.Table 8 compares 50 held-out PushT episodes with identical start/goal pairs; the fixed-terminal controller is described as emphatically horizon-dependent.Source: §C.6 · p. 17, Table 8 · p. 18

D Decoding the Latent Plan: What the Controller Imagines

D.1 The two decoders, and how far each can be trusted

  • TwoRooms decoding reaches PSNR 54.8 dB and SSIM 0.9991 on held-out episodes; PushT is useful for pose and gross geometry, not fine texture.The PushT decoder is 3.7× better than predicting the dataset mean on fresh frames, but remains visibly blurry.Source: §D.1 · p. 18

D.2 Reading the three rows

The three-row diagnostic separates reconstruction error from world-model imagination error, revealing different rollout errors in TwoRooms and PushT. Refinements move imagined endpoints toward goals, while executed episodes illustrate task-specific closed-loop behavior and the limits of decoded claims.

D.2 Reading the three rows

  • TwoRooms imagination stalls at a wall after first-block drift of 0.101, while PushT error grows gradually; the first executed block is accurate.In TwoRooms, drift reaches roughly 1.85 onward and the episode still succeeds at 15.5 pixels against a 16-pixel tolerance; PushT drift grows 0.009 → 0.029 → 0.040 → 0.063 → 0.074.Source: §D.2 · pp. 18–19, Figure 9 · p. 19, Figure 10 · p. 20

D.3 Refinement, made visible

  • Across refinements, imagined terminal states move toward the goal, with most of PushT’s improvement occurring in the first refinement.Predicted distance falls 1.3034 → 1.1884 → 0.3429 → 0.1881 on TwoRooms and 0.1215 → 0.0312 → 0.0309 → 0.0299 on PushT.Source: §D.3 · p. 19, Figure 11 · p. 20, Figure 12 · p. 21

D.4 Executed episodes, end to end

  • The successful episodes show monotone TwoRooms progress but a PushT approach phase in which several replans move the pusher before contact moves the block.The arrival objective must permit apparent lack of progress during those early PushT replans; the displayed episodes illustrate behavior rather than a success rate.Source: §D.4 · p. 21

D.5 The physical trajectory

  • The physical trajectory figure uses trials selected before visual inspection, and decoded claims are limited to TwoRooms and PushT.No quantitative Section 4 result depends on a decoder, which is only a visualization probe.Source: §D.5 · p. 21

E The full algorithm for the architecture:

The controller encodes observations and a goal, iteratively refines a latent action plan using frozen one-step predictions, then executes an initial portion of the final plan. The accompanying episode figures illustrate distinct trial sets and behavior.

E The full algorithm for the architecture:

  • The four PushT executed episodes use sim-generated goals and illustrate behavior, not the held-out expert evaluation success rates.Each row shows ten replans followed by the goal, with final position and orientation error at the right.Source: Figure 13 · p. 22
  • The TwoRooms episodes are held-out trials at m=1, while the physical trajectory shows distance falling from 44.8 to 15.0 pixels.The figure’s trials were fixed before inspection, and its goal markers do not enter the encoder.Source: Figure 14 · p. 22, Figure 15 · p. 23
  • The controller encodes recent observations and the goal, initializes H plan tokens, and executes the first rh action blocks of the final plan.The algorithm repeats after new frames update the history.Source: E The full algorithm for the… · p. 23
  • Each refinement decodes plan tokens into bounded action blocks, rolls them through the frozen predictor, encodes predicted consequences, and applies a learned gated residual correction.The imagined rollout spans H one-step predictions before the plan is corrected.Source: E The full algorithm for the… · p. 23
Loading 2609.13845v1…