LeFlow: Generative Latent Flow Planning for World Models

2026-08-25Computer Vision and Pattern Recognition

Computer Vision and Pattern Recognition
AI summary

The authors propose LeFlow, a new way to plan actions in video-based world models more efficiently. Instead of redoing slow optimization every time a goal changes, LeFlow learns a reusable model that predicts possible future paths directly in a compressed (latent) space. This lets it quickly generate action plans by imagining trajectories and checking them with the world model. Tests show LeFlow is faster and more successful than previous methods on several tasks where agents control pixels to reach goals.

latent world modeltrajectory optimizationlatent spaceconditional trajectory generationinverse dynamics decoderautoregressive rolloutpixel-control benchmarksamortized planningrectified-flow modelgoal-conditioned control
Authors
Hsiang-Wei Huang, Jianxu Shangguan, Junbin Lu, Jenq-Neng Hwang
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.