Control-Geometry Straightening for Sampling-Based Latent Planning

2026-09-28

AI Summary

This study addresses the challenge that latent world models, despite accurate predictions, often yield planning objectives that are difficult to optimize. To this end, we propose a control-geometric straightening method, which is for the first time introduced as a single auxiliary loss within the Joint Embedding Predictive Architecture (JEPA). By aligning the cosine similarities of actions and latent variables, this approach learns planning-friendly representations and theoretically reveals a balancing mechanism between temporal straightening and terminal cost curvature. When integrated with MPPI or CEM sampling planners, experiments demonstrate that our method improves success rates by up to 20% over baselines. Furthermore, with only 128 candidate samples, it significantly reduces the required number of sampling iterations and refinement steps, thereby effectively enhancing overall sampling efficiency.

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.

Problem

• Latent Planning
• Sampling-based Planning
• World Models
• Representation Learning
• Planning Budget

Innovation

• Control-Geometry Straightening
• Latent Planning
• Sampling-based Planners
• World Models
• Representation Learning