arXiv / September 2026World models · JEPA · Control · Planning

Control-Geometry Straighteningfor Sampling-Based Latent Planning

Ziang Fu · Ning Ning

Texas A&M University

Prediction alone doesn’t make a latent space easy to plan through.

TL;DR

Match the geometry of latent changes to the geometry of actions, making good action sequences easier to find under limited planning budgets.

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.

Methodology

One loss to organize control.

Similar actions should induce similar latent changes. Different actions should remain distinguishable—even when the trajectory turns.

For neighboring observations, CGS forms latent differences $\Delta z_t = f_\theta(o_{t+1})-f_\theta(o_t)$ . It matches their pairwise cosine similarities to those of standardized actions within a short history window.

The CGS objective

\mathcal{L}_{\mathrm{CGS}}=\frac{\|VV^\top-UU^\top\|_F^2}{N(N-1)}

Rows of $U=[u_1,\ldots,u_N]^\top$ are unit action directions, with $u_i=a_i/\|a_i\|_2$ ; rows of $V=[v_1,\ldots,v_N]^\top$ are unit latent differences, with $v_i=\Delta z_i/\|\Delta z_i\|_2$ .

Here $a_i$ is a standardized action and $N$ is the number of transitions in the history window. Aligned, orthogonal, and opposing actions provide the corresponding angular targets.

CGS adds this auxiliary loss to LeWM’s prediction and SIGReg objectives. At planning time, the model is frozen and candidate action sequences are refined to bring the predicted endpoint closer to the goal embedding.

Figure 1. Training an encoder and predictor with prediction and CGS losses, then planning by rolling out actions toward an encoded goal. — Fig. 1 Training and planning with CGS. SIGReg is omitted from the diagram.

Understanding the representation

Local alignment.
Balanced planning geometry.

Control becomes legible.

Latent differences follow action-related turns, making changes in control easier to distinguish in the representation.

The benefit survives the horizon.

Over an H-step rollout, CGS yields more balanced endpoint responses in the tested control plane. This helps a planner refine different controllable directions more evenly.

Figure 2. PushT and Cube: actions, latent-transition trajectories for LeWM, TS and CGS, and endpoint-response cosine heatmaps compared with ideal control relations. — Fig. 2 PushT (top) and Cube (bottom). Latent-transition trajectories reveal control structure; endpoint-response heatmaps visualize planning isotropy in a sampled control plane.

A linear-dynamics view

Geometry shapes optimization.

Consider $z_{t+1}=Az_t+Ba_t$ . Under the paper’s fixed-radius and coverage assumptions, the CGS geometry objective encourages three complementary properties.

A\approx I

Less autonomous drift

(A-I)^\top B\approx 0

Less drift–control coupling

B^\top B\approx cI

Balanced action effects

G_H=\Gamma_H^\top\Gamma_H\;\approx\;cH P_H

Terminal-cost curvature becomes more uniform on the active controllable subspace. That balance supports more efficient finite-budget refinement; it does not remove directions that cannot affect the endpoint.

Theory and assumptions in the paper ↗

Planning with CGS

Efficient planning under limited budget

CGS improves sampling-based planning across four environments. The largest gains appear on PushT and Cube; Reacher differences are smaller and depend on the planner and budget.

PushT / MPPI+13.3pp

77.3% vs. 64.0% with LeWM

Cube / MPPI+20.0pp

79.3% vs. 59.3% with LeWM

Same planner.
Same candidate budget.
A different representation.

128 candidates · 30 steps

Goal-reaching success rate (%) at 128 candidates per update. Mean and sample standard deviation over three seeds.
World Model	Planner	PushT	Cube	TwoRooms	Reacher
LeWM	MPPI	64.0 ± 5.3	59.3 ± 4.6	92.0 ± 3.5	65.3 ± 4.2
LeWM	PG-MPPI	72.0 ± 6.0	89.3 ± 4.2	98.0 ± 2.0	62.7 ± 11.0
LeWM	CEM	89.3 ± 6.1	63.3 ± 2.3	76.7 ± 10.1	83.3 ± 8.3
LeWM	PG-CEM	92.0 ± 5.3	90.0 ± 2.0	93.3 ± 3.1	86.7 ± 6.1
LeWM	GD	85.3 ± 5.0	59.3 ± 8.1	41.3 ± 7.0	76.7 ± 4.2
LeWM + TS	MPPI	64.7 ± 8.3 (+0.7)	75.3 ± 5.0 (+16.0)	97.3 ± 3.1 (+5.3)	64.0 ± 2.0 (-1.3)
LeWM + TS	PG-MPPI	69.3 ± 9.0 (-2.7)	92.7 ± 3.1 (+3.3)	98.0 ± 2.0 (+0.0)	66.7 ± 2.3 (+4.0)
LeWM + TS	CEM	86.7 ± 5.0 (-2.7)	73.3 ± 10.3 (+10.0)	97.3 ± 1.2 (+20.7)	81.3 ± 8.1 (-2.0)
LeWM + TS	PG-CEM	90.0 ± 5.3 (-2.0)	90.0 ± 2.0 (+0.0)	100.0 ± 0.0 (+6.7)	78.0 ± 2.0 (-8.7)
LeWM + TS	GD	82.7 ± 5.8 (-2.7)	65.3 ± 4.6 (+6.0)	65.3 ± 9.9 (+24.0)	81.3 ± 5.0 (+4.7)
LeWM + CGS	MPPI	77.3 ± 6.4 (+13.3)	79.3 ± 4.6 (+20.0)	97.3 ± 1.2 (+5.3)	66.7 ± 2.3 (+1.3)
LeWM + CGS	PG-MPPI	80.0 ± 2.0 (+8.0)	96.7 ± 3.1 (+7.3)	100.0 ± 0.0 (+2.0)	69.3 ± 3.1 (+6.7)
LeWM + CGS	CEM	92.0 ± 3.5 (+2.7)	76.7 ± 9.5 (+13.3)	95.3 ± 1.2 (+18.7)	80.7 ± 3.1 (-2.7)
LeWM + CGS	PG-CEM	96.0 ± 0.0 (+4.0)	96.0 ± 4.0 (+6.0)	100.0 ± 0.0 (+6.7)	81.3 ± 3.1 (-5.3)
LeWM + CGS	GD	89.3 ± 6.4 (+4.0)	71.3 ± 4.6 (+12.0)	70.7 ± 3.1 (+29.3)	82.0 ± 5.3 (+5.3)

Three seeds, 50 episodes per seed; mean ± sample SD. Parentheses: change versus LeWM in percentage points; underlined: the two largest gains per task. Bold: best world model per planner and task. Sampling planners use 30 refinement steps; GD uses 100 Adam steps, independent of K. PG denotes a matched action prior.

Planning under a sampling budget

MPPI without an action prior. Shaded bands show ± one sample SD.

LeWM+ TS+ CGS

PushT

Cube

TwoRooms

Reacher

Fig. 3 Success versus the number of candidates per update, with 30 refinement steps.

Planning examples

Two CGS rollouts per environment. Each clip is paired with its goal observation.

PushT

CGS

Contact-rich pushing and rotation.

Cube

CGS

Object transport with a robot arm.

TwoRooms

CGS

Goal-directed navigation through a doorway.

TwoRooms target observation for rollout 1

TwoRooms target observation for rollout 2

Reacher

CGS

Coordinated control of a two-link arm.

Reacher target observation for rollout 1

Reacher target observation for rollout 2

Characterizing CGS · Latent representations and planning

What changes—and why it matters.

Physical motion becomes more explicit.

On PushT, CGS improves both linear and nonlinear motion decoding. Reacher remains strongly decodable nonlinearly, while its state-dependent action effects favor TS for linear organization.

PushT

Agent displacement

Linear decoding · R²

LeWM

0.809

+ TS

0.831

+ CGS

0.903

MLP decoding · R²

LeWM

0.888

+ TS

0.891

+ CGS

0.959

Reacher

Fingertip displacement

Linear decoding · R²

LeWM

0.951

+ TS

0.974

+ CGS

0.899

MLP decoding · R²

LeWM

0.994

+ TS

0.994

+ CGS

0.993

Fig. 4 Held-out motion decoding from latent differences. Higher R² is better.

Compact representations can plan efficiently.

CGS improves both global CLS and spatial Patch features in DINO-WM. LeWM’s jointly learned, compact latent retains strong planning performance at substantially lower rollout cost.

Planning success across DINO CLS, DINO Patch and LeWM with Base, TS and CGS shaping, plus DINO Patch versus LeWM planning times on PushT. — Fig. 5 PushT performance and planner-solve time. DINO-WM models include proprioceptive information; annotated timing ratios are DINO-WM / LeWM.

CEM across representation architectures. Open the full figure for the shared legend. — Fig. 5 PushT performance and planner-solve time. DINO-WM models include proprioceptive information; annotated timing ratios are DINO-WM / LeWM.

Better refinement. A better starting point.

At a fixed candidate budget, CGS reaches high PushT performance in fewer refinement steps. It also improves matched BC-prior rollouts before candidate search begins.

PushT success versus optimization iterations, BC-prior to no-op latent goal distance ratios, and prior-only success across LeWM, TS and CGS. — Fig. 6 Left: CEM/MPPI at K = 128. Middle: prior / no-op distance ratio (lower is better). Right: prior-only planning, K = 1 and I = 1.

CEM refinement steps. Open the full figure for the shared legend. — Fig. 6 Left: CEM/MPPI at K = 128. Middle: prior / no-op distance ratio (lower is better). Right: prior-only planning, K = 1 and I = 1.

BibTeX

@article{fu2026controlgeometry,
  title={Control-Geometry Straightening for Sampling-Based Latent Planning},
  author={Fu, Ziang and Ning, Ning},
  journal={arXiv preprint arXiv:2609.35603},
  year={2026}
}

Prediction alone doesn’t make a latent space easy to plan through.

Abstract

One loss to organize control.

The CGS objective

Local alignment.Balanced planning geometry.

Control becomes legible.

The benefit survives the horizon.

Geometry shapes optimization.

Efficient planning under limited budget

Planning under a sampling budget

PushT

Cube

TwoRooms

Reacher

Planning examples

PushT

Cube

TwoRooms

Reacher

What changes—and why it matters.

Physical motion becomes more explicit.

PushT

Reacher

Compact representations can plan efficiently.

Better refinement. A better starting point.

BibTeX

Local alignment.
Balanced planning geometry.