Latent Geometry Beyond Search: Amortizing Planning in World Models

Motivation

• LeWM predicts efficiently, but its CEM planner requires 9,000 candidate rollouts and 45,000 predictor forward passes per planning step.The authors call the mismatch between cheap prediction and expensive decision-making the planning tax.Source: §1 · p. 1

Contribution

• The paper asks whether planning cost is intrinsic or an artifact of control formulation, and whether regular latent geometry makes action selection a local inverse problem.It replaces online search with a small GC-IDM trained on frozen LeWM latents to predict an action from current state, goal state, and remaining horizon.Source: §1 · pp. 1–2

Results

• Across four benchmark environments, GC-IDM matches or exceeds CEM in seven of eight environment–protocol cells while reducing per-decision planning cost by 100–130×.The results also hold under an episode-level holdout, a 500× CEM compute sweep, and comparisons with other planners.Source: §1 · p. 2

Contribution

• The contributions frame GC-IDM as a single-forward-pass alternative to iterative search and report that it outperforms a broader family of test-time planners.They also report a 100–130× reduction in planning cost.Source: §1 · p. 2

Latent world models and planning over world models

• The paper treats its approach as orthogonal to latent-world-model research: it asks whether the optimizer can be removed because the representation has absorbed its computation.LeWM [Maes et al., 2026] is the backbone, using latent future-embedding prediction and SIGReg regularization [Balestriero and LeCun, 2025].Source: §2 · p. 2

Inverse dynamics for control

• Unlike prior work that uses inverse dynamics as a proposal, auxiliary head, or imitation-learning component, this work makes it the primary planner inside a JEPA latent world model.GLAMOR [Paster et al., 2021] is identified as the closest precursor, predicting action sequences from pixels with a recurrent model and learned action prior.Source: §2 · p. 2
• GC-IDM predicts one action per step with a feedforward MLP on frozen pretrained JEPA latents, then replans from a fresh observation.It conditions directly on the terminal goal embedding and remaining horizon, without an explicit state-prediction stage inside the controller.Source: §2 · pp. 2–3

Problem formulation

• The controller encodes current and goal observations into latent states, then seeks actions that minimize their latent goal mismatch.The goal-matching cost is typically induced by Euclidean distance in latent space.Source: §3 · p. 3

Search-based planning in latent space

• With sufficiently regular latent geometry, the optimizer can be replaced by a learned inverse map that recovers the first action from current latent, goal latent, and horizon.Figure 1 illustrates Push-T embeddings evolving from diffuse clouds into a smooth ordered path aligned with task progress.Source: §3 · p. 3, Figure 1 · p. 4

Latent world model

• A latent world model combines an observation encoder with a predictor of one-step latent dynamics under actions.Its training objective jointly optimizes the encoder and predictor, with a SIGReg regularizer encouraging smooth, well-distributed embeddings and trading off against predictive accuracy.Source: §3 · p. 3

Search-based planning in latent space

• Standard latent-space planning minimizes terminal goal distance by repeatedly rolling out candidate action sequences through the learned predictor.LeWM’s default CEM configuration uses 9,000 candidate rollouts and 45,000 predictor forward passes per planning step.Source: §3 · p. 3

4.1 Overview

• The method’s motivation is that decoding well-posed local transitions can work even when linearly interpolating between distant latent states yields poor plans.A pairwise IDM achieved R2 = 0.993 for oracle action reconstruction but planned poorly, indicating the difficulty was constructing a valid latent path.Source: §4.1 · p. 4
• GC-IDM predicts the next action from the current latent, goal latent, and remaining horizon, avoiding imagined latent trajectories and online search.The controller uses actual encoder outputs from observations and amortizes the local inverse map offline.Source: §4.1 · p. 5

4.2 Model Architecture and Details

• GC-IDM is a roughly 1.5M-parameter, three-layer MLP trained from scratch in about 20 minutes per environment on a single GPU.It concatenates current and goal embeddings; the remaining-step count is normalized, sinusoidally encoded, and applied through AdaLN-Zero modulation.Source: §4.2 · p. 5
• Training samples a current latent, a future latent at a uniformly sampled horizon, and the corresponding action from the offline demonstration dataset.It reuses the LeWM training dataset and requires no additional environment interaction; Figure 2 depicts the frozen-embedding training and single-pass planning pipeline.Source: §4.2 · p. 5
• Table 1 compares GC-IDM and CEM across four LeWM benchmarks at sample sizes n=200 and n=50, reporting success and timing measures.The caption states that results are means ± standard deviations over three seeds and that the reported speedup is amortized over CEM’s 25-action commit window.Source: Table 1 · p. 6

4.1 Overview

• At test time, GC-IDM encodes the goal once, re-encodes the current observation after each action, and replans using the remaining horizon.The procedure is closed-loop, with no inner optimization or open-loop commitment to a predicted latent trajectory.Source: §4.1 · p. 5, §4.2 · p. 6

5.1 Experimental Setup

• The evaluation uses goal-conditioned image control in Two-Room navigation, Push-T contact-rich manipulation, OGBench-Cube object manipulation, and Reacher continuous reaching, with n=50 and n=200 protocols.It also includes a strict episode-level holdout and supplementary compute and ablation studies.Source: §5.1 · p. 6

5.2 Implementation Details

• GC-IDM is a horizon-conditioned MLP trained on frozen LeWM embeddings, adding about 20 minutes of per-environment fitting in exchange for one encoder pass and one MLP pass per control step.CEM repeatedly invokes the latent predictor during its optimization loop.Source: §5.2 · pp. 6–7

5.3 Main Results

• Seven of eight environment–protocol cells favor or match GC-IDM over CEM, with 100–130× lower planning cost; Push-T at n=50 is the exception.There, CEM is modestly better in mean success but more than two orders of magnitude slower.Source: §5.3 · p. 7

5.4 Comparison with MPPI, iCEM, and Gradient-Based Planners

• At matched defaults, GC-IDM has the highest success in every (n, env) cell and costs 29–134× less per plan call than the broader solver family.The comparison includes CEM, MPPI, iCEM, and GradientSolver; the best sampling baseline varies by environment.Source: §5.4 · p. 8, Figure 4 · p. 8

5.5 Ablation Studies

• Across a 500× CEM compute sweep, no configuration is both faster and more successful than GC-IDM on any environment.On Push-T, the best CEM setting reaches 90% at 11.4 s per episode, exceeding GC-IDM’s 84.2% at 35× the planning cost.Source: §5.5 · p. 8
• Two-Room stays at 100% across goal offsets and Reacher and Cube remain nearly flat, while Push-T degrades with distance; short horizons hurt, with gains largely saturating by Hmax=50.The horizon result is at the 50-step evaluation budget.Source: §5.5 · p. 8

5.6 Trajectory Quality: Smoothness and Directness

• GC-IDM produces 15–36× smoother actions across four environments and approaches the goal more directly in latent space on three, with Cube near tied at 69% versus 71%.On Two-Room, the matched rollout reports 92% latent monotonicity for GC-IDM versus 73% for CEM, and action jerk of 0.055 versus 1.962.Source: Figure 3 · p. 7, §5.6 · p. 9, Table 4 · p. 9

6 Conclusion and Future Work

• A 1.5M-parameter GC-IDM trained in about 20 minutes per environment replaces the 9,000-rollout CEM loop with one forward pass per step, matching or exceeding CEM in seven of eight cells at 100–130× lower cost.It is trained on frozen LeWM embeddings.Source: §6 · p. 9
• The success and cost advantage also extends to MPPI, iCEM, and gradient-based optimization, supporting learned inference in sufficiently well-structured latent spaces as an alternative to online optimization.This is the authors’ interpretation of the reported results.Source: §6 · p. 9
• The regime’s reach remains to be tested with stronger irreversibility, longer-horizon contact structures, higher-dimensional actions, and different latent architectures.The authors also propose comparing isotropic and non-isotropic latent spaces.Source: §6 · p. 9

Mechanism 1: No train/test distribution shift

• GC-IDM uses real ViT-encoded observations during both training and testing, avoiding the distribution shift that drives pairwise IDM to require extreme noise augmentation.The noise ablation reports that every environment prefers σ=0.Source: A Why GC-IDM Works: Mechanism Analysis · p. 12

Mechanism 2: Closed-loop correction via re-encoding

• GC-IDM re-encodes every environment step, so a single-step prediction error is overwritten by a fresh observation rather than accumulated across committed actions.Default CEM commits 25 raw environment actions before re-encoding; over a 50-step episode, that is 2 re-encodings for CEM versus 50 for GC-IDM.Source: A Why GC-IDM Works: Mechanism Analysis · p. 12

Mechanism 3: Directionally correct predictions accumulate

• Even with low pointwise R2, the IDM can learn the average direction toward the goal, and per-step replanning accumulates directionally correct predictions into successful trajectories.On Two-Room, 92% of GC-IDM episodes show monotonically decreasing latent distance to the goal at every step.Source: A Why GC-IDM Works: Mechanism Analysis · p. 12

Combined mechanism

• Together, the mechanisms replace trajectory enumeration with a directed action step followed by fresh encoding; the analysis predicts noise insensitivity, a need for horizon supervision, and little benefit from extra capacity.The paper associates these predictions with Tables E, F, and H, respectively.Source: A Why GC-IDM Works: Mechanism Analysis · p. 12

• The framework defines latent transitions by applying the environment dynamics between encoded states, with the inverse state mapping understood only locally on the data manifold.Source: B Theoretical Analysis: Isotropy and… · p. 13
• The analysis assumes approximately whitened, isotropic latent statistics and locally differentiable, bi-Lipschitz encoding with a full-column-rank action Jacobian near demonstrations.The local latent density is also idealized as approximately Gaussian in the visited region.Source: B Theoretical Analysis: Isotropy and… · p. 13
• Under these assumptions, nearby latent displacements satisfy δz ≈ Bt δa, yielding a continuous local inverse when actions induce distinguishable latent motion.The minimum-norm local solution uses the Moore–Penrose pseudoinverse.Source: B Theoretical Analysis: Isotropy and… · p. 13
• When the encoder is locally near-isometric, its Jacobian adds little ill-conditioning, making the induced inverse problem easier to approximate; distributional isotropy alone does not guarantee this.The argument connects this condition to SIGReg discouraging collapsed or disproportionately amplified latent directions.Source: B Theoretical Analysis: Isotropy and… · p. 13

C More Qualitative Comparisons

• GC-IDM produces smoother action sequences in Two-Room and Reacher, while contact-rich Push-T and OGBench-Cube show greater CEM step-to-step variance and occasional GC-IDM undershoot.The reported Push-T failure typically undershoots goal-pose rotation rather than diverging entirely; the figure contrasts one action per step with CEM candidate rollouts and replanning.Source: C More Qualitative Comparisons · p. 14
• No CEM configuration was both faster and more successful than GC-IDM across the 500× compute sweep.The sweep covered 12 configurations and 60–30,000 rollouts; the shaded faster-and-more-successful region was empty in every panel.Source: C More Qualitative Comparisons · pp. 14–15
• Across four matched Cube episodes, GC-IDM succeeded in 65, 91, 155, and 192 steps while CEM failed at 200 steps in each.Source: C More Qualitative Comparisons · p. 18
• Across four matched Reacher episodes, GC-IDM succeeded in 161–185 steps while CEM failed at 200 steps in each.The listed GC-IDM episode times were 5.6–5.7 seconds, versus 45.0–45.4 seconds for CEM.Source: C More Qualitative Comparisons · p. 19

D Pilot Study on Pairwise IDM

• The pairwise IDM is presented as the empirical precursor motivating the goal-conditioned formulation, alongside oracle, noise, and closed-loop failure analyses.Source: D Pilot Study on Pairwise IDM · p. 20

D.1 Pairwise IDM

• The pairwise IDM maps consecutive latent states to their intervening action; frameskip-1 training reaches R2 = 0.993, whereas frameskip-5 yields R2 ≈0 on Two-Room.The authors attribute the ambiguity to many distinct multi-action sequences connecting the same two latent states.Source: §D.1 · p. 20

D.2 Noisy Conditioning

• Gaussian input noise is added during IDM training to address the mismatch between ground-truth embeddings at training and interpolated or predictor-generated embeddings at test time.Source: §D.2 · p. 20

D.3 Oracle Experiment and Noise Ablation

• Ground-truth-pair action reconstruction reaches R2=0.993 at σ=0, but the best closed-loop success is 34% with training noise σ=1.5.Reconstruction degrades gracefully up to test-time noise σ≈0.5 and collapses beyond σ=1; the authors identify trajectory generation, rather than action decoding, as the bottleneck.Source: §D.3 · p. 20

D.4 Closed-Loop Performance and Failure Mode

• 48% is the pairwise IDM’s ceiling on Two-Room, versus CEM’s 87%; refinement and noisy candidates still leave a large gap.Refinement raises success from 34% to 42%; adding 10 noisy candidates brings it to 48%, while iterations beyond three yield diminishing returns.Source: §D.4 · p. 21

E.1 Data Efficiency

• Two-Room and Reacher approach saturation with just 5% of training data, while Push-T success continues to rise with more data.OGBench-Cube improves smoothly toward near-saturation with the full training set.Source: §E.1 · p. 22

E.2 Noise Augmentation is Unnecessary for GC-IDM

• GC-IDM performs best without injected noise; on Push-T, raising σ from 0 to 0.5 lowers success from 85.5% to 67.0%.No environment prefers σ > 0; unlike pairwise IDM, GC-IDM uses real encoder outputs at both training and test time.Source: §E.2 · p. 22

E.3 Role of the Horizon Embedding

• Removing horizon supervision reduces success by an average of 42 percentage points, and widening the network does not recover the loss.Hmax=1 is equivalent to single-step goal-conditioned behavior cloning in latent space and does not train on distant goals.Source: §E.3 · pp. 22–23

E.4 Full Solver-Family Comparison

• GC-IDM is the highest-success method in seven of eight (n, environment) cells; only Push-T at n=50 favors CEM, with overlapping seed noise.At n=200, the Push-T ranking reverses; among sampling baselines, none closes the gap to GC-IDM.Source: §E.4 · p. 23

E.5 Architecture Hyperparameter Ablation

• On Push-T, depth improves success from 70.5% with one layer to 87.5% with five, while width has negligible effect across an 8× range.Two-Room and Reacher are saturated across configurations; the default three-layer model is within 2.5 percentage points of the best observed setting.Source: §E.5 · p. 23

E.6 Noise Schedule

• Uniform noise schedules do not recover performance lost to noise injection: Push-T remains 4–7 percentage points below its 85.5% no-noise baseline.Two-Room, Reacher, and OGBench-Cube are invariant to schedule and magnitude.Source: §E.6 · p. 23

E.8 Held-Out Validation

• On held-out episodes, success is 100.0% for Two-Room, 84.8% for Push-T, 98.5% for OGBench-Cube, and 99.8% for Reacher.These rates are close to the in-distribution results, and the authors report held-out performance as statistically indistinguishable from them.Source: §E.8 · pp. 24–25
• Held-out wall-clock speedups range from 2.2× on OGBench-Cube to 40× on Two-Room.Source: §E.8 · p. 25

E.7 Evaluation Budget and Long-Horizon Control

• By 25 steps, Two-Room reaches 100% success and OGBench-Cube peaks at 99.5%; Push-T peaks at 93% and falls to 75% by 100 steps.Push-T’s decline is consistent with horizon-signal clamping beyond Hmax=50, a setting not seen during training; Reacher reaches 99.5% by 50 steps and holds.Source: §E.7 · p. 24, §E.8 · p. 26

IDM architecture

• Pairwise and goal-conditioned IDMs use three 512-dimensional hidden layers, LayerNorm, GELU, and 10% dropout.The output head is initialized with σ=0.01 to start near zero.Source: F More Implementation Details · p. 27

Training hyperparameters

• IDMs are trained for 50 epochs with AdamW at a 10^-3 learning rate, batch size 1024, cosine annealing, and gradient clipping at norm 1.0.Weight decay is 10^-4, and the learning rate anneals to 1/100 of its initial value.Source: F More Implementation Details · p. 27

Goal-conditioned IDM specifics

• GC-IDM uses Hmax=50 and samples goal frames uniformly from 1–50 steps ahead within the same episode.The horizon embedding is a 64-dimensional sinusoidal encoding.Source: F More Implementation Details · p. 27

Embedding extraction

• Frames are encoded in batches of 256 through the frozen LeWM encoder with ImageNet normalization; HDF5 data are streamed to limit memory use.DataParallel is used when multiple GPUs are available.Source: F More Implementation Details · p. 27

Evaluation and environment configurations

• Evaluation uses 200 parallel environments, random starts, goals offset by 25 steps, and a 50-step budget; planning time includes encoding and inference.The protocol matches LeWM’s per-environment evaluation configuration, and the CEM solver uses horizon, receding horizon, and action block values of 5.Source: F More Implementation Details · p. 27, Table L · p. 27

G More Limitations and Future Works

• The study focuses on LeWM and a benchmark suite with 2–5 DoF; higher-dimensional actions, other world-model representations, and severely irreversible tasks remain open extensions.For severe irreversibility, the authors suggest GC-IDM as a fast proposal with CEM as an occasional corrector.Source: G More Limitations and Future Works · p. 28

NeurIPS Paper Checklist

• The paper reports no theoretical contributions and says its appendices document hyperparameters, evaluation details, solver setup, and available LeWM checkpoints and datasets.The checklist reports mean ± standard deviation across three training seeds for the main and solver-family comparisons.Source: NeurIPS Paper Checklist · p. 29
• The checklist says the code will be released upon acceptance, while the LeWM checkpoints and datasets are publicly available.Source: NeurIPS Paper Checklist · p. 29