Control-Geometry Straightening for Sampling-Based Latent Planning
Abstract
Joint-embedding predictive architectures enable planning with latent world models, but accurate transition prediction alone does not ensure that the planning objective is easy to optimize. We introduce Control-Geometry Straightening (CGS), a single auxiliary loss that learns planner-friendly representations by directly straightening control geometry for sampling-efficient planning. CGS matches pairwise cosine similarities among actions to those among corresponding latent differences only using local transitions from pixel-action pairs. The loss can be applied across world-model architectures using end-to-end learned or pretrained representations. Under linear-dynamics, our theoretical analysis connects this objective to temporal straightening and more balanced terminal-cost curvature across the full planning horizon, yielding finite-budget guarantees for MPPI, local contraction results for CEM, and convergence bounds for gradient descent. Across four control environments and multiple planners, CGS improves planning with fewer sampled candidates and refinement steps, achieving success-rate gains up to 20 and 12.6 percentage points over LeWorldModel (LeWM) and its temporal-straightening variant (LeWM+TS), respectively, with sampling-based planners using 128 candidates per update. Probes, comparisons with DINO-WM architecture, and planner-side ablations clarify how latent motion organization, state dependence, and dynamical context shape planning behavior. Straightening control geometry thus makes good action sequences easier to find under limited planning budgets.
Related
- Latent Geometry and the Emergence of Lorentzian Time in Matrix Quantum Mechanics
- Stochastic Geometry Meets Optimization: Data-Driven Cellular Planning
- Latent Annealed Posterior Sampling with a Variance-Preserving Bridge for Sparsely Observed Random Fields
- Geometry and observability of latent texture-coherence diffusions in coherent sea-clutter radar observations
- Latent-Space No-Arbitrage Geometry of Generative Models for Implied Volatility Surfaces
BibTeX
@misc{fu2026controlgeometry,
title = {{Control-Geometry Straightening for Sampling-Based Latent Planning}},
author = {Ziang Fu and Ning Ning},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2609.35603},
eprint = {2609.35603},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.35603},
}
RIS
TY - UNPB
TI - Control-Geometry Straightening for Sampling-Based Latent Planning
AU - Ziang Fu
AU - Ning Ning
PY - 2026
DA - 2026/10/07
PB - arXiv
DO - 10.48550/arxiv.2609.35603
UR - https://arxiv.org/abs/2609.35603
AB - Joint-embedding predictive architectures enable planning with latent world models, but accurate transition prediction alone does not ensure that the planning objective is easy to optimize. We introduce Control-Geometry Straightening (CGS), a single auxiliary loss that learns planner-friendly representations by directly straightening control geometry for sampling-efficient planning. CGS matches pairwise cosine similarities among actions to those among corresponding latent differences only using local transitions from pixel-action pairs. The loss can be applied across world-model architectures using end-to-end learned or pretrained representations. Under linear-dynamics, our theoretical analysis connects this objective to temporal straightening and more balanced terminal-cost curvature across the full planning horizon, yielding finite-budget guarantees for MPPI, local contraction results for CEM, and convergence bounds for gradient descent. Across four control environments and multiple planners, CGS improves planning with fewer sampled candidates and refinement steps, achieving success-rate gains up to 20 and 12.6 percentage points over LeWorldModel (LeWM) and its temporal-straightening variant (LeWM+TS), respectively, with sampling-based planners using 128 candidates per update. Probes, comparisons with DINO-WM architecture, and planner-side ablations clarify how latent motion organization, state dependence, and dynamical context shape planning behavior. Straightening control geometry thus makes good action sequences easier to find under limited planning budgets.
ER -