PaperlayerView exact revision on arXiv

Source-linked AI summary

LeFlow: Generative Latent Flow Planning for World Models

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

arXiv:2608.24855v1cs.CV

TL;DR

Existing latent-world-model planners repeatedly solve each state–goal query through online black-box optimization. LeFlow amortizes planning with a reusable latent trajectory prior, decoding and rollout-verifying candidates, and reports higher success with roughly an order-of-magnitude less planning time across four benchmarks.

  • Problem

    Existing planners solve every state–goal query from scratch with repeated online trajectory optimization, reusing no structural planning knowledge across queries.

  • Method

    LeFlow generates conditional latent paths between encoded current and goal observations, decodes latent transitions into action chunks, and verifies candidates through frozen-world-model rollouts.

  • Results

    LeFlow achieves the highest mean success on all four benchmarks, including 100.0 ± 0.0% on TwoRoom, while cutting planning time by roughly an order of magnitude.

  • Takeaways & Limitations

    Planning on latent world models can be amortized into a reusable trajectory prior rather than re-solved from scratch for each query.

  • Takeaways & Limitations

    LeFlow currently operates at a fixed short horizon, while longer horizons increase generation dimensionality and accumulated rollout error.

Abstract

from arXiv · show

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.

1. Introduction

Latent world models provide strong compact dynamics predictors, but planning still repeatedly searches action sequences online from scratch. LeFlow instead amortizes planning as latent trajectory generation, decodes actions afterward, and verifies candidates with the frozen world model.

  • Background: Latent world models encode observations into compact representations and predict future latent states conditioned on actions.LeWorldModel is presented as a stable JEPA that trains end-to-end from raw pixels.
  • Result: LeFlow achieves better success with an order-of-magnitude reduction in planning time under the Figure 1 evaluation protocol.The comparison uses LeWM+CEM and LeFlow over 50 episodes with horizon H=5.
  • Problem: Online CEM planning treats the world model as a black-box simulator and repeats thousands of queries for every state–goal replanning step.Planning experience from one query is not reused for the next.
  • Approach: LeFlow learns a reusable latent trajectory prior offline instead of re-solving each planning problem online.The prior operates directly in the frozen world model’s latent space.
  • Approach: LeFlow generates conditional latent paths, decodes transitions into executable action chunks, and reranks candidates using frozen-world-model rollouts.This separates planning from control and evaluates the actual rollout distance to the goal.

2. Related Work

Prior work uses latent world models with online black-box optimization, inverse-dynamics models, or generative trajectory planners. LeFlow differs by generating explicit latent paths inside a frozen world model, decoding actions locally, and reranking candidates by rollout.

  • Latent world models: Latent world models predict action-conditioned dynamics in compressed spaces, while test-time planning commonly uses CEM or MPPI as black-box optimizers.LeWM is used as LeFlow’s frozen backbone.
  • Amortized planning and inverse dynamics: Inverse-dynamics approaches predict actions responsible for latent transitions, including goal-conditioned models that combine inverse dynamics with learned action priors.LeFlow instead separates latent trajectory generation from local action decoding.
  • Amortized planning and inverse dynamics: LeFlow generates explicit multi-step latent trajectories with a flow model and decodes actions locally from each transition.It operates in the latent space of a frozen pretrained world model rather than jointly learning representation and planner.
  • Generative models for trajectory planning: Diffuser-style methods model joint state–action trajectories, whereas LeFlow generates latent state paths and decodes actions afterward via inverse dynamics.The related work motivates generative modeling because multiple valid trajectories can connect the same start and goal.

3. Method

LeFlow replaces repeated action-space optimization with conditional latent-path generation on a frozen LeWM, then decodes and verifies candidate trajectories. Its modular design separates multi-step latent planning from local action recovery and rollout-based feasibility selection.

  • 3. Method: LeFlow generates reusable goal-conditioned latent trajectories instead of repeatedly searching directly over action sequences.The approach exploits the world model’s compact latent geometry and delegates action realization to a separate decoder.
  • 3. Method: A frozen LeWM supplies latent representations, autoregressive dynamics, and rollout-based verification, while LeFlow trains only its planner and decoder.LeWM’s encoder provides the latent anchors and its predictor evaluates whether generated proposals follow realizable dynamics.
  • 3. Method: The rectified-flow planner generates only the H−1 interior states between clamped start and goal latent anchors.Fixing endpoints injects goal conditioning directly into the path structure and concentrates capacity on how the latent state evolves between anchors.
  • 3. Method: An inverse dynamics decoder converts each latent transition into an executable action chunk using the current state, next state, and displacement.This separates global trajectory-shape reasoning from the local problem of realizing one latent step.
  • 3. Method: Generated paths may leave the controllable latent manifold, so candidates are decoded, rolled out through frozen LeWM, ranked by terminal goal distance, and the best is executed.Reranking selects among samples without modifying the flow model’s distribution.
  • 3. Method: Training combines flow matching, inverse-dynamics learning, and a consistency constraint that steers generated transitions toward predictor-reachable dynamics.The consistency loss shapes the training distribution, while rollout reranking selects the best proposal at inference.

4. Experiments

Across four goal-conditioned pixel-control benchmarks, LeFlow combines latent-path generation, inverse dynamics, and frozen-world-model rollout verification to replace online action-space optimization. It achieves higher success than LeWM+CEM while reducing planning time by roughly an order of magnitude.

  • Main results: LeFlow reaches 100.0 ± 0.0% on TwoRoom, 95.2 ± 3.0% on PushT, 86.8 ± 4.8% on Reacher, and 100.0 ± 0.0% on OGBench-Cube.These results use a single planner design and attain the highest mean success on every benchmark.
  • Main results: +18.0, +5.9, +18.8, and +26.7 points over LeWM+CEM are reported on TwoRoom, PushT, Reacher, and OGBench-Cube, respectively.The comparison supports replacing iterative online action-space optimization with amortized latent planning.
  • Planning efficiency: LeFlow is 14.4×, 11.4×, 11.7×, and 4.5× faster than LeWM+CEM on TwoRoom, PushT, Reacher, and OGBench-Cube, respectively.Both methods use the same frozen backbone and protocol; end-to-end timing includes shared environment stepping and rendering.
  • Planner design ablation: Removing either latent-space planning or generative modeling degrades success, with LeFlow best on the non-saturated PushT and Reacher benchmarks.Latent-space planning exceeds direct action generation by +3.0 on PushT and +4.5 on Reacher, while generative modeling exceeds deterministic latent paths by +1.0 and +4.5.
  • Rollout reranking ablation: Rollout reranking improves Reacher by +10.5 points (77.0→87.5) and PushT by +2.5 points over executing the first sampled candidate.Candidates are scored using their actual frozen-LeWM rollout distance to the goal.
  • Consistency-loss ablation: Performance is stable across tested λcons values, with λcons=0.1 used in main models and a marginal but consistent +2.0-point gain on Reacher.At λcons=1.0, performance slightly degrades on both benchmarks, consistent with reduced candidate diversity.

5. Conclusion and Future Work

LeFlow amortizes planning on a frozen LeWM by generating latent paths, decoding actions, and verifying candidates through frozen-model rollouts. Across four benchmarks it improves success while cutting planning time by roughly an order of magnitude, but currently remains limited to short fixed horizons.

  • Conclusion: LeFlow generates goal-conditioned latent paths, decodes them into action chunks through inverse dynamics, and verifies candidates by frozen-model rollout.This replaces repeated online optimization with one proposal-and-reranking pass.
  • Conclusion: Across four pixel-control benchmarks, LeFlow improves success while reducing planning time by roughly an order of magnitude.The method trades offline planner training for lower online planning cost.
  • Future work: LeFlow currently operates at a fixed short horizon, where longer horizons increase generation dimensionality and accumulated frozen-predictor rollout error.The paper identifies hierarchical chains of shorter-horizon segments and stronger world models as future directions.

A. Benchmark Descriptions

LeFlow is evaluated on four continuous-action, goal-conditioned pixel-control benchmarks spanning 2D navigation, 2D manipulation, sparse robotic reaching, and visually complex 3D manipulation.

  • Benchmark scope: The evaluation covers four goal-conditioned pixel-control benchmarks, all using continuous action spaces.Goals are pixel observations from the same episode as the start and are sampled a fixed number of steps ahead to ensure reachability.
  • 2D navigation: TwoRoom requires navigating through a single doorway because a wall blocks the direct path to a target in the other room.The agent must first move toward the door and then redirect toward the target.
  • 2D manipulation: PushT requires multi-step pushing to match a T-shaped block’s target position and orientation without grasping.The benchmark uses 20,000 expert episodes, with a 50-step evaluation budget and goals sampled 25 steps ahead.
  • Robotic reaching: Reacher tests coordinated torque control for a two-joint planar arm reaching a randomly placed target in a visually sparse scene.Its sparse visuals make the task primarily a test of dynamics modeling rather than visual complexity.
  • 3D manipulation: OGBench-Cube evaluates 3D manipulation in which a robotic arm picks up and places a cube at a target position.The single-cube setting uses 10,000 episodes of 200 steps, with realistic lighting and six-degree-of-freedom pose variation.

B. Held-Out Generalization

LeFlow retains near-identical success on episode-level held-out queries across all four benchmarks, indicating that its planner generalizes beyond episodes used for planner training.

  • Evaluation protocol: LeFlow is retrained on 80% of episodes and evaluated exclusively on start–goal queries from the held-out 20%.The frozen LeWM backbone remains identical across the split.
  • Results: Held-out success rates are statistically indistinguishable from in-distribution performance on every benchmark.The comparison uses the same five-seed, 50-episode protocol as the main evaluation.
  • Results: PushT held-out success is 97.2 ± 1.1% versus 95.2 ± 3.0% in-distribution.The held-out result is slightly higher, while the reported difference remains within the overall noise range.
  • Results: Reacher held-out success is 86.0 ± 5.7% versus 86.8 ± 4.8% in-distribution, while OGBench-Cube is 100.0 ± 0.0% in both settings.Together, these results show negligible held-out gaps on the remaining benchmarks.

C. Architecture and Training Details

LeFlow combines a Transformer-based rectified-flow planner with an MLP inverse-dynamics decoder, trained jointly with fixed optimization and inference settings.

  • Latent-path planner: The velocity model vψ is a 4-layer pre-norm Transformer with hidden dimension d=512, 8 attention heads, and feedforward width 4d=2048.The model uses GELU activations.
  • Latent-path planner: Flow time τ is projected to R512 and combined with learned positional embeddings, supporting up to 19 interior steps with Hmax=20.The time embedding is sinusoidal and is processed through two linear layers with SiLU activation.
  • Inverse dynamics: The inverse-dynamics decoder is a 3-layer MLP whose input concatenates zt, zt+1, and the displacement zt+1−zt.Its hidden dimension is 512, with LayerNorm and GELU after each hidden layer and no dropout.
  • Training and inference: Both modules are trained jointly with AdamW for 10 epochs using lr=10^-4, weight_decay=10^-4, batch size 128, and gradient clipping at norm 1.0.The loss weights are λflow=1.0, λinv=1.0, and λcons=0.1.
  • Training and inference: Inference samples N=64 paths and integrates the flow for 16 Euler steps.These settings define the fixed sampling and integration budget used by the planner.

D. Horizon-scaling Results

PushT performance deteriorates sharply as the planner horizon increases, while runtime rises modestly; the authors attribute this to both method and benchmark limitations and keep the main focus on amortizing short-horizon optimization.

  • Scaling results: PushT success decreases from 94% at H=5 to 32% at H=10 and 6% at H=20.Each latent transition is decoded into a five-step action block, and each horizon is trained separately.
  • Scaling results: Planning runtime increases from 28.00 seconds at H=5 to 34.73 seconds at H=10.All settings use 50 evaluation cases, a 100-step execution budget, five-step receding horizon, 64 proposals, and 16 flow steps.
  • Interpretation: The reported degradation reflects accumulated flow, inverse-dynamics, and rollout-prediction errors over longer paths.The authors also note that PushT and other official LeWM benchmarks primarily target short-term control rather than long-horizon scaling.
  • Scope: The paper does not aim to extend LeWM to long-horizon planning tasks, identifying hierarchical planning and longer-term benchmarks as future directions.The stated objective is to amortize online trajectory optimization through learned latent-path generation.

E. Generalization to DINO-WM

LeFlow transfers its latent-path planning and verification design to DINO-WM on PushT, with success improving as training data and proposal count increase. At 512 proposals, it surpasses matched CEM while using substantially fewer rollouts and less time, although evidence remains limited to one task and a small bottleneck-training set.

  • Representation and setup: LeFlow applied to DINO-WM by compressing its 256 × 384 DINOv2 patch grid and combining the resulting visual vector with a 10-dimensional proprioceptive embedding.The proprioceptive input includes agent position and velocity, while the pushed object’s pose and orientation remain visually inferred.
  • Representation and setup: The planner generates an H=5 compact-space path, decodes each transition into five low-level actions, and uses the frozen DINO-WM transition model for candidate rollout and reranking.The five-action decoding matches DINO-WM’s frame skip.
  • Results: 38%, 60%, 68%, and 78% success were achieved with 1k, 5k, 10k, and 15k trajectories, respectively, using 64 proposals.The success rate increases monotonically with planner-training data under the fixed 64-proposal setting.
  • Results: 86% success at 512 proposals surpassed DINO-WM+CEM at 84%, while mean state distance decreased from 25.41 to 15.55.The 15k-trajectory LeFlow model produced these results when only the inference proposal count was increased from 64 to 512.
  • Results: LeFlow used 17.6× fewer candidate rollouts and reduced end-to-end time from 37:09 to 2:32 across 50 cases, a 14.7× reduction versus CEM.LeFlow reranked 512 candidates in one pass, whereas CEM used 30 refinement iterations over 300 candidates.
  • Scope: The transfer evidence is limited to PushT, and the spatial bottleneck was trained on only 1k trajectories, so universal backbone independence is not established.The authors characterize this as evidence of transfer to a second backbone rather than universal backbone independence.
Loading 2608.24855v1…