LeFlow: Generative Latent Flow Planning for World Models

Hsiang-Wei Huang; Jianxu Shangguan; Junbin Lu; Jenq-Neng Hwang

cs.CVAug 25, 2026

LeFlow: Generative Latent Flow Planning for World Models

Authors: Hsiang-Wei Huang, Jianxu Shangguan, Junbin Lu, Jenq-Neng Hwang

Organizations: University of Washington, United States

Abstract

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.

Figures & tables

Figure 1: LeFlow plans an order of magnitude faster than CEM at better success. End-to-end evaluation time (left) and success rate (right) for LeWM+CEM and LeFlow under the same protocol (

H{=}5

50

episodes). Replacing online black-box search with an amortized latent trajectory prior removes the per-step optimization loop while preserving task performance.

Figure 2: LeFlow architecture. A frozen LeWM encoder maps the current and goal observations to latent anchors $z_{\text{start}}$ and $z_{\text{goal}}$ . A conditional rectified-flow planner generates the interior of a latent path between them; an inverse dynamics decoder converts each latent transition into an action chunk. Only the flow planner and inverse dynamics decoder are trained; LeWM encoder is frozen throughout.

Figure 3: Rollout reranking projects generative proposals onto the controllable latent manifold. The flow model samples

N

candidate latent paths (dashed) with endpoints clamped to

z_{\text{start}}

and

z_{\text{goal}}

, but generated interiors may pass through latent states unreachable by any admissible action sequence. Each candidate is decoded into actions and rolled out autoregressively through the frozen LeWM predictor (solid), producing an actual terminal latent

\hat{z}_{H}^{(i)}

. Candidates are ranked by rollout distance to the goal (Eq. 5 ); the best is executed. This step grounds amortized proposals in the predictor’s dynamics without modifying the flow model.

Method	TwoRoom	PushT	Reacher	OGBench-Cube
Random	0.0	2.0	10.0	48.0
GCBC [ 7 ]	100.0	75.0	–	84.0
GCIVL [ 19 ]	100.0	33.0	–	56.0
GCIQL [ 13 ]	100.0	20.0	–	64.0
PLDM [ 25 ]	97.0	78.0	78.0	65.0
DINO-WM [ 28 ]	100.0	74.0	79.0	86.0

Table 1: Main benchmark results: success rate (%) on four goal-conditioned pixel-control tasks. Baseline numbers (goal-conditioned BC and offline RL – GCBC, GCIVL, GCIQL; the JEPA world model PLDM; and the foundation-encoder world model DINO-WM) are as reported by Maes et al. [18] ; “–” marks benchmarks for which a baseline is not reported. LeFlow is the mean

\pm

std over five independent evaluation runs (seeds

42

–

46

) under the LeWM codebase-default protocol (

H{=}5

50

episodes). Best per column in bold .

Benchmark	Method	Success (%) $\uparrow$	Eval Time (s) $\downarrow$	Speedup $\uparrow$
TwoRoom	LeWM + CEM	82.0 $\,\pm\,$ 2.0	224.78	1.0 $\times$
TwoRoom	LeFlow	100.0 $\,\pm\,$ 0.0	15.58	14.4 $\times$
PushT	LeWM + CEM	89.3 $\,\pm\,$ 6.4	198.92	1.0 $\times$
PushT	LeFlow	95.2 $\,\pm\,$ 3.0	17.42	11.4 $\times$
Reacher	LeWM + CEM	68.0 $\,\pm\,$ 9.2	326.01	1.0 $\times$
Reacher	LeFlow	86.8 $\,\pm\,$ 4.8	27.67	11.8 $\times$

Table 2: Planning efficiency: LeWM+CEM vs. LeFlow. Both methods use the same frozen LeWM backbone and the same evaluation protocol (

H{=}5

50

episodes). Success rates are reported as in Table 1 (LeFlow: mean over five runs, seeds

42

–

46

; LeWM+CEM: mean

\pm

std over three seeds under stable_worldmodel defaults); eval time for CEM and LeFlow are both five-run mean. LeFlow exceeds CEM success on every benchmark while reducing end-to-end planning time by roughly an order of magnitude. Speedup is the ratio of CEM to mean LeFlow eval time.

	LeFlow	Action-Flow		Det.-Latent
Benchmark	(%)	(%)	$\Delta$	(%)	$\Delta$
PushT	96.5	93.5	$+3.0$	95.5	$+1.0$
Reacher	87.5	83.0	$+4.5$	83.0	$+4.5$

Table 3: LeFlow vs. two planner ablations. Action-Flow generates action chunks directly (no latent states or inverse dynamics) and uses the same rollout reranking as LeFlow; Det.-Latent replaces the rectified-flow planner with a deterministic regressor (single candidate, no reranking).

H{=}5

200

episodes;

\Delta

is LeFlow minus the variant (positive favors LeFlow). Non-saturated benchmarks only. See Sec. 4.3 .

Benchmark	Full Rerank (%)	No Rerank (%)	$\Delta$
PushT	96.5	94.0	$+2.5$
Reacher	87.5	77.0	$+10.5$

Table 4: Rollout reranking ablation. “Full” scores candidates by frozen-LeWM rollout distance to the goal; “No rerank” executes the first sampled candidate.

H{=}5

200

episodes;

\Delta

is Full minus No-rerank. Non-saturated benchmarks only. See Sec. 4.4 .

	$\lambda_{\text{cons}}{=}0.0$	$\lambda_{\text{cons}}{=}0.1$ (main)	$\lambda_{\text{cons}}{=}1.0$
PushT	96.5	96.5	95.0
Reacher	85.5	87.5	84.5

Table 5: Consistency-loss ablation. Success rate (%) as the consistency weight

\lambda_{\text{cons}}

varies;

\lambda_{\text{cons}}{=}0.1

(main) is used by all main models.

H{=}5

200

episodes. See Sec. 4.5 .

Figure 4: LeFlow rollout visualizations. Each row shows a single successful episode: start observation (left) through intermediate executed steps to goal achievement (right). Across navigation (TwoRoom), pushing (PushT), reaching (Reacher), and 3D manipulation (OGBench-Cube), LeFlow produces consistent, goal-directed behavior.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

	TwoRoom	PushT	Reacher	OGBench-Cube
LeFlow (in-dist.)	$100.0\pm 0.0$	$95.2\pm 3.0$	$86.8\pm 4.8$	$100.0\pm 0.0$
LeFlow (held-out)	$100.0\pm 0.0$	$97.2\pm 1.1$	$86.0\pm 5.7$	$100.0\pm 0.0$

Appendix

Table 6: Held-out generalization. LeFlow retrained on 80% of episodes and evaluated exclusively on the held-out 20%. Results are mean

\pm

std over five seeds; in-distribution numbers (from Table 1 ) shown for reference. Differences are within noise on every benchmark.

H	Success (%)	Time (s)
5	94.0	28.00
10	32.0	30.57
20	6.0	34.73

Appendix

Table 7: PushT horizon scaling. Each planner is trained separately at its reported horizon. All settings use 50 evaluation cases sampled with seed 42, a 100-step budget, and an action block of five.

Method (data)	Proposals	SR (%)	Distance $\downarrow$	Time (m:s)
LeFlow (1k)	64	38.0	38.60	0:43
LeFlow (5k)	64	60.0	23.88	0:43
LeFlow (10k)	64	68.0	26.53	0:34
LeFlow (15k)	64	78.0	19.40	0:34
LeFlow (15k)	512	86.0	15.55	2:32
DINO-WM+CEM	$300{\times}30$	84.0	25.41	37:09

Appendix

Table 8: LeFlow on a frozen DINO-WM backbone (PushT). All methods use the same 50 evaluation cases. LeFlow rows use the same frozen spatial encoder and visual-plus-proprio representation. Time is end-to-end job time for all 50 cases; state distance is DINO-WM’s final-state metric.

Latent-Space Planning

See all 31 papers

Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent Dynamics Modeling World Model-Based Planning
Xi’an Jiaotong University
Latent Geometry Beyond Search: Amortizing Planning in World Models
May 9, 2026Hoang Nguyen, Xiaohao Xu, Xiaonan HuangWorld Model-Based Planning Model-Based Planning
Department of Robotics, University of Michigan, Ann Arbor
LePlanner: An Iterative Amortized Controller For World Models
Sep 12, 2026Saksham Bansal, Om Naphade, Chayan Aggarwal +1World Model-Based Planning Model-Based Planning
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 +1

Jun 24, 2026J/K move · Enter open · S save

Fast LeWorldModel

Yuntian Gao, Xiangyu Xu

Xi’an Jiaotong University

Latent Dynamics Modeling World Model-Based Planning Joint-Embedding Predictive Architecture Action-Conditioned World Models Efficient Inference for World Action Models Model-Based Planning+1

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.

Abstract

Latent-Space Planning

Explore similar work

Fast LeWorldModel

Latent Geometry Beyond Search: Amortizing Planning in World Models

LePlanner: An Iterative Amortized Controller For World Models