---
abstract: |
  Learning good representations is essential for latent planning with world models. While pretrained visual encoders produce strong semantic visual features, they are not tailored to planning and contain information irrelevant---or even detrimental---to planning. Inspired by the perceptual straightening hypothesis in human visual processing, we introduce temporal straightening to improve representation learning for latent planning. Using a curvature regularizer that encourages locally straightened latent trajectories, we jointly learn an encoder and a predictor of a Joint-Embedding Predictive Architecture (JEPA) world model. We show that reducing curvature this way makes the Euclidean distance in latent space a better proxy for the geodesic distance and improves the conditioning of the planning objective. We demonstrate empirically that temporal straightening makes gradient-based planning more stable and yields significantly higher success rates across a suite of goal-reaching tasks. Our code is available at <https://agenticlearning.ai/temporal-straightening/>.
author:
- |
  Ying Wang$^{1}$, Oumayma Bounou$^{1}$, Gaoyue Zhou$^{1}$, Randall Balestriero$^{2}$, Tim G. J. Rudner$^{3}$,\
  Yann LeCun$^{1}$$^{*}$, and Mengye Ren$^{1}$$^{*}$\
  $^{1}$New York University $^{2}$Brown University $^{3}$University of Toronto\
  Correspondence to `{yw3076,yann.lecun,mengye}@nyu.edu`\
  <https://agenticlearning.ai/temporal-straightening>
bibliography:
- main.bib
---

# Temporal Straightening for Latent Planning

## Abstract

Learning good representations is essential for latent planning with world models. While pretrained visual encoders produce strong semantic visual features, they are not tailored to planning and contain information irrelevant---or even detrimental---to planning. Inspired by the perceptual straightening hypothesis in human visual processing, we introduce temporal straightening to improve representation learning for latent planning. Using a curvature regularizer that encourages locally straightened latent trajectories, we jointly learn an encoder and a predictor of a Joint-Embedding Predictive Architecture (JEPA) world model. We show that reducing curvature this way makes the Euclidean distance in latent space a better proxy for the geodesic distance and improves the conditioning of the planning objective. We demonstrate empirically that temporal straightening makes gradient-based planning more stable and yields significantly higher success rates across a suite of goal-reaching tasks. Our code is available at <https://agenticlearning.ai/temporal-straightening/>.

# Introduction

Latent world models offer a compelling solution for planning due to better efficiency and generalization [@truck-backer-upper; @dyna; @worldmodel; @Hafner2020Dream; @hafner2022dreamerv2; @hafner2023dreamerv3; @hansen2022temporal; @hansen2023td]. They compress high-dimensional observations into compact latent representations, learn predictive dynamics in that latent space, and enable imaginary rollouts for action optimization. Compared to operating directly in pixel or state space, the latent abstraction reduces dimensionality and ignores noise, making dynamics learning more efficient. At test time, planning is typically posed as optimizing an action sequence by rolling the model forward and minimizing a cost function between the goal and the predicted states in the latent space.

<figure id="img:teaser">
<img src="images/wall_teaser.png" />
<figcaption>Latent trajectories encoded by a pretrained visual encoder are usually highly curved, increasing the difficulty of prediction and planning. We learn a representation space where feasible trajectories are straighter to facilitate latent planning. </figcaption>
</figure>

In practice, however, optimization in the learned latent space remains challenging. The induced planning objective is typically highly non-convex, potentially causing gradient-based optimizers to struggle. As a result, many successful practices [@hafner2019planet; @hansen2023td; @dinowm; @sobal2025learning; @terver2026drivessuccessphysicalplanning] rely on search-based methods such as CEM [@rubinstein1997optimization] or MPPI [@mppi], which achieve competitive performance but introduce a substantial compute burden and latency. Moreover, commonly used goal cost metrics based on Euclidean distance can be misleading if the embedding space is not properly regularized. In particular, when latent trajectories are highly curved, straight-line distances in embedding space misrepresent the geodesic distance along feasible transitions. These challenges call for better representations that facilitate latent planning.

```{=latex}
\begin{figure*}

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=\textwidth]{images/pusht_exp_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=\textwidth]{images/umaze_exp_pca_mse.pdf}
    \end{subfigure}
\caption{Latent trajectories before vs.\ after straightening. The upper PushT example is a rotation and the bottom UMaze example shows the agent traveling from the top-left to the top-right, with the star denoting the target. Straightening yields less curved and smoother trajectories, and makes Euclidean distance a more faithful proxy for geodesic progress towards the goal. More examples are in~\Cref{app:pca}.}

\label{img:pca_mse}
\end{figure*}
```
What is a \`\`good" representation for latent planning? Although general-purpose visual pretraining provides powerful semantic-aware features, it is not tailored to the dynamics of the environment and often retains plenty of planning-irrelevant low-level details. We argue that planning could benefit from representations that are (i) sufficient for predicting dynamics but without task-irrelevant information and (ii) properly regularized so that embedding distances reflect the geodesic distance and gradient-based optimization is reliable. With such representations, we can exploit the differentiability of latent world models and enable efficient gradient-based planning, bypassing the need for computationally expensive search-based methods.

Inspired by the perceptual straightening hypothesis in human vision [@perceptual_straightening], which posits that visual systems transform complex natural videos into straighter internal representations, we introduce a simple approach to straighten latent trajectories for planning. Concretely, we jointly learn an encoder and a predictor of a Joint-Embedding Predictive Architecture (JEPA) world model, while imposing regularization on the curvature of latent trajectories during training. We find that the JEPA prediction objective alone induces *implicit straightening* to some extent, and introducing an *explicit* curvature regularizer further strengthens and stabilizes this effect. The resulting encoded trajectories are significantly straighter, with Euclidean distances better aligned with geodesic distances (`\Cref{img:pca_mse}`{=latex}). We prove that reducing curvature improves convergence of gradient-based planners, and observe superior empirical gains across a suite of goal-reaching tasks: open-loop planning success improves by 20--60% and MPC by 20--30% with a simple gradient-based planner.

```{=latex}
\begin{figure*}[t]\center
\includegraphics[width=0.9\textwidth]{images/architecture.pdf}
\caption{During training, we minimize the prediction loss between the predicted embedding $\hat{z}_{t}^t$ and the target ${z}_{t}^t$ with stop-grad in the target branch, and minimize the local curvature of embeddings. During planning, we roll out for the horizon $T$ using the trained predictor and select optimal actions that minimize the cost between the predicted terminal state and the goal in the embedding space.
}
\label{img:train}
\end{figure*}
```
# Related Work

While early visual world models directly predict in pixel spaces and use generated images for control [@oh2015actionconditionalvideopredictionusing; @finn2017deep; @ebert2018visual; @unipi], an increasing number of recent works first encode high-dimensional sensory inputs into compact latent representations and plan in the resulting latent space. Learning representations is central to these latent world models.

To obtain meaningful representations for world modeling, prior methods add reconstruction-based objectives when training the encoder along with the predictor [@watter2015embedcontrollocallylinear; @zhang2019solardeepstructuredrepresentations; @Levine2020pcc; @worldmodel; @hafner2019planet; @Hafner2020Dream; @iris; @transformer100k]. However, these reconstruction objectives overemphasize low-level visual details that are unnecessary for planning and may fail to capture task-relevant information. More recent approaches decouple perception from dynamics by leveraging strong pretrained visual encoders [@nair2022r3muniversalvisualrepresentation; @dinowm; @bar2025navigation; @dexwm; @bai2025wholebodyconditionedegocentricvideo; @assran2025vjepa2selfsupervisedvideo]. Closest to our setup, DINO-WM [@dinowm] trains task-agnostic predictors and plans directly in frozen DINOv2 [@oquab2024dinov2] feature space. While DINOv2 features provide high-quality visual representations, they are not optimized for planning and may lead to planning objectives that are challenging to optimize. In this work, we improve the representation space for planning by introducing a curvature regularization during world model training.

Joint-Embedding Predictive Architecture (JEPA) emerges as a promising paradigm for world models by learning representations through prediction [@jepa; @vjepa; @assran2025vjepa2selfsupervisedvideo]. It aims to capture predictable structure without retaining unpredictable low-level details, making it more effective and efficient than reconstruction objectives [@vanassel2025jointembeddingvsreconstruction]. This paradigm has been shown to be effective for predictive modeling and planning, with training from scratch on offline simulator data [@sobal2025learning] and large-scale real-world video pretraining followed by action-conditioned post-training for robotic planning [@assran2025vjepa2selfsupervisedvideo]. Our work also belongs to the JEPA family and focuses on learning better representations.

Temporal Contrastive Learning [@sermanet2018timecontrastivenetworksselfsupervisedlearning; @tclr; @eysenbach2024inference; @yang2025memory] is also a popular paradigm for learning representations that can reflect the temporal relationships. It encourages temporally close frames to have similar embeddings while distant frames have more dissimilar embeddings through InfoNCE loss [@radford2021learningtransferablevisualmodels]. However, how to choose positive and negative samples requires careful tuning and this objective might push away geodesically close states if suboptimal trajectories are used. Instead, our regularization-based method does not require negatives and applies *local* straightening without requiring expert trajectories.

Motivated by the perceptual straightening hypothesis in human vision [@perceptual_straightening], some prior works have examined implicit straightening in pretrained visual encoders [@harrington2023exploring; @interno2025aigenerated] or used it as an objective to obtain robust video models [@niu2024learning; @bagad2025chirality]. Relatedly, early work on learning linearized video representations also regularizes the curvature of latent trajectories [@goroshin2015learninglinearizeuncertainty]. Implicit straightening is also observed in autoregressive language models optimized for next-word prediction [@hosseini2023large; @hosseini2026context]. To the best of our knowledge, however, none of these prior works have studied the impact of straightening on world modeling and planning. We further discuss the connection of our work to a broader literature in control and learning plannable representations in `\Cref{sec:related_cont}`{=latex}.

# Temporal Straightening

We consider control tasks with high-dimensional observations $o_t \in \mathbb{R}^{n_o}$ of an agent interacting with its environment using actions $a_t \in \mathbb{R}^{n_a}$. Our goal is to learn a world model that maps observations to a latent space and models the dynamics in this space, which we use for latent planning.

In this section, we first outline the architecture of our world model, then define the training objectives with a novel geometric regularization that straightens latent trajectories.

## World Model

Our world model predicts future states in a learned latent space and consists of three components: a sensory encoder, an action encoder, and a predictor.

#### Sensory encoder.

The sensory encoder $\mathcal{E}^s_\phi$ maps raw observations $o_t$ into latent representations $$\begin{aligned}
z_t \in \mathbb{R}^{d} = \mathcal{E}^s_\phi(o_t).
\end{aligned}$$

The sensory encoder can be any function that maps observations to latent representations. For visual observations, the encoder may preserve spatial structure or collapse it into a global vector representation.

#### Action encoder.

Each action $a_t \in \mathbb{R}^{n_a}$ is mapped to a latent action embedding via $$\mathcal{E}^a_\psi: \mathbb{R}^{n_a} \to \mathbb{R}^{d_a}.$$

#### Predictor.

The predictor `\mbox{$f_\theta: \mathbb{R}^{K \times d} \times \mathbb{R}^{K \times d_a} \to \mathbb{R}^{d}$}`{=latex} models transitions in the latent space. Given a history of $K$ past latent states and actions, it predicts the next latent state $$\hat{z}_{t} = f_\theta \left( \{z_i\}_{i=t-K}^{t-1}, \{\mathcal{E}^a_\psi(a_i)\}_{i=t-K}^{t-1} \right).$$

## Straightening Latent Trajectories

We seek to straighten the latent space induced by the sensory encoder $\mathcal{E}^s_\phi$ by penalizing the curvature along trajectories. Let $z_t, z_{t+1}$, and $z_{t+2}$ be three consecutive latent representations obtained by encoding observations $o_t, o_{t+1}$, and $o_{t+2}$ using $\mathcal{E}^s_\phi$ . We define approximate latent velocity vectors $$v_t = z_{t+1} - z_t, \quad v_{t+1} = z_{t+2} - z_{t+1},$$ and seek to minimize the angle between them, or equivalently maximize their cosine similarity $$\label{eq:cossim}
\mathcal{C}= \frac{v_t \cdot v_{t+1}}{||v_t||_2 \cdot ||v_{t+1}||_2}.$$

## Training Objective.

The parameters $\phi, \psi$ and $\theta$ of the world model components $\mathcal{E}^s_\phi$, $\mathcal{E}^a_\psi$ and $f_\theta$ are trained jointly to minimize prediction error and enforce straightened trajectories.

#### Prediction objective.

We minimize the MSE between the predicted and target latent states $\hat{z}_{t+1}$ and $z_{t+1}$:

$$\mathcal{L}_{pred} = || \hat{z}_{t+1} - \operatorname{sg}(z_{t+1}) ||_2^2 \\ \text{ },
\label{eq:pred}$$ where $\mathrm{sg}$ denotes the stop-gradient operation to prevent collapse of the latent space.

#### Straightening objective.

We minimize trajectory curvatures by minimizing the negative cosine similarity $$\mathcal{L}_{curv} = 1 - \mathcal{C}.
\label{eq:curvature}$$ This straightening loss can be applied to any differentiable sensory encoder, either in isolation or jointly with the prediction objective.

#### Overall objective.

The total training objective combines prediction and straightening as

$$\mathcal{L}_{total} = \mathcal{L}_{pred} + \lambda \mathcal{L}_{curv},$$

where $\lambda \geq 0$ controls the strength of the straightening.

#### Collapse prevention.

Since our encoder is trainable, the model is likely to produce degenerate solutions in which all latent representations collapse to a constant. Common anti-collapse strategies can be regularization-based [@bardes2022vicregvarianceinvariancecovarianceregularizationselfsupervised; @zhu2024variancecovarianceregularizationimprovesrepresentation; @balestriero2025lejepa; @kuang2026rectified], contrastive-based [@SimCLR; @moco], and stop-gradient-based [@chen2020simsiam; @byol]. Our curvature regularizer is *orthogonal to these anti-collapse methods* and can be combined with any of them. We use stop-grad for major experiments due to its simplicity and efficiency, as it does not require negative samples or introduce new hyperparameters. We apply stop-gradient to the target latent in the prediction loss `\eqref{eq:pred}`{=latex} to prevent the gradients from $\mathcal{L}_{pred}$ from being backpropagated through the target branch. Although a collapsing solution is still possible in theory, stop-grad has been shown to be effective in self-supervised vision learning [@chen2020simsiam], and also effective in our experiments.

<figure id="img:loss_landscape">
<figure id="fig:wall">
<img src="images/landscape_baseline.png" />
<figcaption>DINOv2</figcaption>
</figure>
<figure id="fig:pointmaze">
<img src="images/landscape_ours.png" />
<figcaption>Straightened</figcaption>
</figure>
<figcaption> Action-Space Loss Landscape. We pick one test sample from PushT with a planning horizon of 25 steps. For each coordinate <span class="math inline">(<em>a</em><sub><em>x</em></sub>, <em>a</em><sub><em>y</em></sub>)</span> in the grid, we fix the first action and optimize the remaining actions in the planning horizon to minimize the terminal goal cost. The heatmap represents the minimum attainable loss for each initial action choice, with darker colors indicating lower loss. The loss landscape is closer to being convex after straightening. </figcaption>
</figure>

# Planning with Straightened Dynamics {#sec:theory}

In this section, we present a theoretical analysis on the effect of straightening in the case of a linear dynamical system and show that straightened latent dynamics lead to better convergence in gradient-based planning.

We consider a goal-reaching task where we optimize an action sequence $\mathbf a=(a_0,\dots,a_{K-1})\in\mathbb R^{K\times d_a}$ over a horizon $K$ to reach a target latent goal $z_g$. For simplicity, we use the mean-squared terminal error, $$\label{eq:obj}
\mathcal L(\mathbf a)=\|z_K-z_g\|_2^2,
\qquad
z_K=\Phi(\mathbf a),$$ where $\Phi$ denotes unrolling the learned latent dynamics from a fixed initial state $z_0$.

::: assumption
**Assumption 1** (Linear latent dynamics). `\label{as:linear}`{=latex} For analysis, we consider linear latent dynamics $$\label{eq:linear}
f: (z_t, a_t) \to Az_t+Ba_t,
\qquad
\text{s.t.} \quad
z_{t+1}=Az_t+Ba_t,$$ where $A\in\mathbb R^{d\times d}$ and $B\in\mathbb R^{d\times d_a}$. We first state results for $d_a=d$ and $B$ invertible; see Remark `\ref{rem:lowdim}`{=latex} for $d_a<d$.
:::

::: definition
**Definition 2** ($\varepsilon$-straight transition). `\label{def:eps}`{=latex} Under the linear dynamics, we call $f$ *$\varepsilon$-straight* if $$\label{eq:eps_def}
\|A-I\|_2 \leq \varepsilon.$$ The term \`\`straight" reflects that, as $\varepsilon$ tends to $0$, the dynamics of $f$ approach those of the reference function $g: (z_t, a_t) \to z_t+Ba_t$, where the state evolves linearly along a straight trajectory modified only by the control input. We are primarily interested in the regime where $\varepsilon$ is small.
:::

::: remark
*Remark 3* (Cosine similarity as a practical proxy). `\label{rem:cos_proxy}`{=latex} In practice, we regularize temporal straightness using the cosine similarity between consecutive latent velocities `\eqref{eq:cossim}`{=latex}. Under mild bounded-variation assumptions on velocity magnitudes and smooth actions, a large cosine similarity implies that $(A-I)$ is small along visited velocity directions. Detailed proofs are in `\Cref{app:theory_cos}`{=latex}.
:::

::: theorem
**Theorem 4** (Conditioning of the planning Hessian). *`\label{thm:cond}`{=latex} Under Assumption `\ref{as:linear}`{=latex} with $d_a=d$ and $B$ invertible, unrolling `\eqref{eq:linear}`{=latex} yields $$z_K=A^{K}z_0+\sum_{t=0}^{K-1}A^{K-1-t}Ba_t,$$ so $z_K$ is affine in $\mathbf a$ and the planning Hessian is $$H:=\nabla_{\mathbf a}^2\mathcal L(\mathbf a)=2J_\Phi^\top J_\Phi\succeq 0,$$ where $J_\Phi=\bigl[\,A^{K-1}B,\; A^{K-2}B,\; \cdots,\; B\,\bigr]
\in\mathbb R^{d\times Kd}.$ Let $\mathcal W_K:=J_\Phi J_\Phi^\top=\sum_{k=0}^{K-1}A^kBB^\top(A^\top)^k$ be the finite-horizon controllability Gramian [@kailath1980linear; @sontag1998mathematical; @chen1999linear]. Then the *effective* condition number $\kappa_{\mathrm{eff}}(H):=\sigma_{\max}(H)/\sigma_{\min}^+(H)$ satisfies $$\label{eq:kappaA}
\begin{aligned}
\kappa_{\mathrm{eff}}(H)=\kappa(\mathcal W_K)
&\le
\kappa(B)^2\,
\frac{\sum_{k=0}^{K-1}\sigma_{\max}(A)^{2k}}
     {\sum_{k=0}^{K-1}\sigma_{\min}(A)^{2k}} \\
&\le
\kappa(B)^2\,\kappa(A)^{2(K-1)},
\end{aligned}$$ where $\kappa(A):=\sigma_{\max}(A)/\sigma_{\min}(A)$. Moreover, if the transition is $\varepsilon$-straight with $\varepsilon=\|A-I\|_2<1$, then $$\label{eq:eps_bound}
\kappa_{\mathrm{eff}}(H)
\;\le\;
\kappa(B)^2\left(\frac{1+\varepsilon}{1-\varepsilon}\right)^{2(K-1)}.$$ For $\varepsilon\le\tfrac12$, this gives $\kappa_{\mathrm{eff}}(H)
\;\le\;
\kappa(B)^2 e^{6\varepsilon K}$.*

*Proofs are in `\Cref{app:theory_hessian}`{=latex}.*
:::

::: remark
*Remark 5* (Low-dimensional actions). `\label{rem:lowdim}`{=latex} When $d_a<d$, $B$ is not invertible and $\mathcal W_K$ (hence $H$) may be singular outside the controllable subspace. Theorem `\ref{thm:cond}`{=latex} holds on $\mathcal S_K=\mathrm{range}(\mathcal W_K)$ when $\kappa_{\mathrm{eff}}$ is computed using $\sigma_{\min}^+(\mathcal W_K)$; see `\Cref{app:theory_hessian}`{=latex}.
:::

Theorem `\ref{thm:cond}`{=latex} shows that $\varepsilon$-straight transitions control the condition number of the planning Hessian: when $\varepsilon=\|A-I\|_2$ is small, the Gramian remains better conditioned, yielding $\kappa_{\mathrm{eff}}(H)$ that grows slowly with the horizon. Since the planning objective is quadratic with Hessian $H\succeq 0$, gradient descent converges linearly at a rate controlled by the condition number, so the improved bounds on $\kappa_{\mathrm{eff}}(H)$ translate to faster optimization in practice. For nonlinear predictors $z_{t+1}=f_\theta(z_t,a_t)$, analogous guarantees require controlling products of state-dependent Jacobians and higher-order terms, which can be an exciting future work direction.

Empirically, we observe that straightening yields a loss landscape with reduced non-convexity under nonlinear dynamics (`\Cref{img:loss_landscape}`{=latex}). In the next section, we show that it improves gradient-based planning.

# Experiments

To test the effectiveness of the proposed method, we evaluate planning on four environments: Wall [@dinowm; @sobal2025learning], PointMaze UMaze and a more complicated medium maze [@fu2020d4rl], and PushT [@chi2024pusht]. We compare against the baseline DINO-WM [@dinowm], which builds on frozen DINOv2 spatial features or CLS tokens. Following DINO-WM's setup, we use a frameskip of 5 for all environments. Details on the environments and experiments are in `\Cref{app:datasets,app:experiments}`{=latex}.

## Architecture details {#subsec:architecture-details}

Here, we describe the encoder and predictor architectures used to instantiate our world model.

#### Visual encoder.

We consider two encoder setups for the visual encoder $\mathcal{E}^s_\phi$:

-   A **frozen pretrained** visual backbone with a lightweight projector: We use DINOv2 [@oquab2024dinov2] as the backbone.[^1] Given an observation, the backbone produces spatial features $e_t \in \mathbb{R}^{M \times D}$. We add a trainable lightweight CNN projector $\mathcal{P}_\phi$ on top of the backbone, leading to $$z_t^v = \mathcal{P}_\phi(e_t) \in \mathbb{R}^{m_v \times d_v},$$ where we usually choose $m_v \le M$ and $d_v \le D$. The projector may reduce spatial resolution (pooling/striding), channel dimension, or both, encouraging abstraction and reducing computation.

-   A ResNet [@He2015DeepRL] trained **from scratch**, producing features $z_t^v \in \mathbb{R}^{m_v \times d_v}$ directly.

#### Predictor.

We use a ViT [@dosovitskiy2021vit] as the dynamics predictor $f_\theta$. When available, proprioceptive states $p_t \in \mathbb{R}^{n_p}$ are encoded via $\mathcal{E}^p_\xi: \mathbb{R}^{n_p} \to \mathbb{R}^{d_{p}}$ and concatenated with each visual spatial feature. To condition on actions, we concatenate the action embeddings $z_t^a=\mathcal{E}^a_\psi(a_t)\in\mathbb{R}^{d_a}$ with the visual and proprioceptive embeddings before passing them to the predictor. We apply a temporal causal attention mask so tokens at time $t$ attend only to frames $\{t-K,\dots,t-1\}$, enabling frame-level autoregressive prediction.

#### Cosine similarity computation.

The straightening loss (Eq. `\eqref{eq:cossim}`{=latex}) is only applied on visual latents $z_t^v$. Different implementations depend on whether latent representations preserve spatial structure:

-   **Global features** ($n_v = 1$): Compute the cosine similarity directly between vectors.

-   **Spatial features** ($n_v > 1$): We consider four variants: (i) compute the cosine similarity per-patch and average across patches; (ii) flatten all spatial features and then compute the cosine similarity; (iii) average-pool the spatial features before cosine similarity; (iv) use a learnable aggregation head to aggregate spatial features before cosine similarity. We use (iv) in the main experiments and ablate these choices in `\Cref{app:straightening}`{=latex}.

```{=latex}
\begin{figure*}[t]\center

\includegraphics[width=\linewidth]{images/curvature_bars_scatter_hw14.pdf}
\caption{Latent Curvature and Open-Loop GD Success Rate for Different Encoders. Higher cosine similarity indicates lower curvature. Here, we compare models with spatial features and report the average patch-wise cosine similarity. Given the same type of encoder, reduced curvature generally leads to higher success rates.
}
\label{img:cos}
\end{figure*}
```
```{=latex}
\begin{figure*}[t]\center
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/umaze_dino_cls.pdf}
        \caption{DINOv2 CLS embedding}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/resnet/umaze_aggmlpcos1e-1_resnetplus_dim8_hw14/20_agg.pdf}
        \caption{Straightened (agg head)}
        \label{img:heatmap_umaze_proj_agg}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/resnet/umaze_aggmlpcos1e-1_resnetplus_dim8_hw14/20_patch.pdf}
        \caption{Straightened (spatial features)}
        \label{img:heatmap_umaze_proj_patch}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/umaze_geo8_gt.pdf}
        \caption{Ground-Truth using A-star}
    \end{subfigure}
\caption{Distance heatmaps of PointMaze (blue indicates small values, and red indicates large values). The yellow star represents the target, and we compute the Euclidean distance between its embedding and those of all other states in the maze. \Cref{img:heatmap_umaze_proj_agg,img:heatmap_umaze_proj_patch} use ResNet with output features $z \in \mathbb{R}^{14\times14\times8}$, trained with straightening regularization. After straightening, the latent distance accurately reflects the minimum number of steps required to reach the target.
}
\label{img:heatmap}
\end{figure*}
```
## How good is the embedding space? {#sec:analysis_emb}

We first inspect the learned embedding space before comparing the downstream planning performance. We measure latent trajectory curvatures and latent Euclidean distances to understand the impact of straightening. For interpretability, we train a VAE [@kingma2013vae; @Oord2017NeuralDR] decoder with a reconstruction loss, detaching latents to stop gradients from the encoder and predictor.

We find that (i) *implicit straightening* can happen in JEPA world models when training the encoder using the prediction loss alone; (ii) adding explicit curvature regularization further strengthens and stabilizes the straightening effect; (iii) straightening encourages the latent Euclidean distance to better align with the geodesic distance; (iv) near-perfect reconstruction can be attained with a very low feature dimensionality.

```{=latex}
\begin{figure*}

    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_wall_dino.png}
        \caption{Wall: DINO}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_wall_ours.png}
        \caption{Wall: Ours}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_umaze_dino.png}
        \caption{UMaze: DINO}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_umaze_ours.png}
        \caption{UMaze: Ours}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_medium_dino.png}
        \caption{Medium: DINO}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.15\textwidth}

        \includegraphics[width=\textwidth]{images/gd_exp_medium_ours.png}
        \caption{Medium: Ours}
    \end{subfigure}

\caption{Comparison of Open-loop GD Planning. The star denotes the target. For each subfigure, the upper row shows the overlaid rendered images from the simulator by executing the actions, and the bottom shows imaginary rollouts (with a frameskip of 5) decoded using a trained decoder. GD planners can easily get stuck with pretrained DINOv2 features, while straightening significantly increases success. More examples of open-loop planning are in~\Cref{app:traj}.}
\label{img:example_rollout}

\end{figure*}
```
#### Reduced curvature.

In `\Cref{img:cos}`{=latex}, we compare the curvature of test latent trajectories by computing the cosine similarity of the difference in adjacent frames as in `\Cref{eq:curvature}`{=latex}. We also visualize the latent trajectories using PCA as shown in `\Cref{img:pca_mse,app:pca}`{=latex}.

The pretrained DINOv2 embedding space is highly curved as shown in the PCA plots and reflected by the low cosine similarities. The embedding space generally becomes straighter after training even without explicit straightening regularization. We attribute this *implicit straightening* to the JEPA objective: it favors representations whose temporal evolution is easy to predict, so training pressure reduces abrupt directional changes in the latent trajectory. With the *explicit straightening* regularization, the curvature of the embedding space is effectively reduced further. We observe that training a ResNet encoder from scratch generally yields lower curvatures than training a projector on top of a frozen pretrained backbone, likely because it offers greater representational flexibility to adapt the geometry to the dynamics.

When straightening is applied to the aggregation head, the curvature of the aggregated features is significantly reduced while the underlying spatial features are not forced to be overly straightened. For example, PushT has more complex object motions and the patch-wise cosine similarity is unable to faithfully capture the global state changes. The introduction of an aggregation head increases the flexibility of representation learning and generally leads to better planning performance (see `\Cref{app:straightening}`{=latex}). We thus use this implementation for the main experiments.

#### Faithful distance.

Although DINOv2 is a strong visual encoder for various downstream vision tasks, it is not optimized for planning and control. As shown in `\Cref{img:pca_mse}`{=latex}, MSE (which is equal to the squared Euclidean distance) between pretrained DINOv2 features does not reflect the progress of moving towards the target. To better understand the limitation of DINOv2, we visualize the Euclidean distance between the embedding of a target state and all other states in the maze in `\Cref{img:heatmap}`{=latex}. We also compare with ground-truth geodesic distance maps, computed using the A-star search algorithm on the grid of the maze. More heatmaps from different environments and encoders are in `\Cref{app:heatmap}`{=latex}.

Straightening results in a distance heatmap that closely aligns with the geodesic distance. Notably, the model is only trained on suboptimal, non-expert trajectories. Yet, it does not simply memorize the inefficient paths from the training data; instead, it learns to approximate the minimum number of steps required to transition between states. We also find that the spatial features and aggregated global features capture different levels of distance information. The spatial features preserve local geometry and thus yield fine-grained, locally discriminative distance variations, whereas global features provide a smoother, more coherent long-range signal that better reflects long-horizon distance-to-goal trends.

#### Sufficient information.

To examine whether or not these projected features preserve sufficient information for planning, we train a decoder to reconstruct images from latents. The decoder is solely for interpretability purposes and is detached from the world model via stop-gradient. Note that perfect reconstruction is not required, since planning only depends on task-relevant information. However, in our visually simple environments, even aggressively compressed features reconstruct the observations with high fidelity, as shown in `\Cref{img:example_rollout}`{=latex}. This indicates that the resulting features retain sufficient planning-relevant information.

```{=latex}
\begin{table*}[t]

  \caption{Goal-reaching Success Rate of 50 Test Samples (\%) using the GD planner. Values are mean $\pm$ std over three data sampling seeds. The best values are \textbf{bold}. The shaded rows are ours while the rest is DINO-WM~\cite{dinowm}.}
  \label{tab:main_results}
  \setlength{\tabcolsep}{5pt}
  \renewcommand{\arraystretch}{1.08}
  \resizebox{\linewidth}{!}{%
  \begin{tabular}{lcc@{}|@{}cc@{}cc@{}cc@{}cc}
    \toprule
    Encoder & Config &
    Wall &
    PointMaze -- UMaze &
    PointMaze -- Medium &
    PushT \\
    \cmidrule(lr){2-3}\cmidrule(lr){4-5}\cmidrule(lr){6-7}\cmidrule(lr){8-9}\cmidrule(lr){10-11}
     & dim & $\mathcal{L}_{curv}$
     & Open-loop & MPC & Open-loop & MPC & Open-loop & MPC & Open-loop & MPC \\
    \midrule
    DINOv2 (CLS) & $1 \times 384$ & \xmark & \pmstat{28.67}{12.68} & \pmstat{66.67}{10.50} & \pmstat{25.33}{0.94} & \pmstat{82.67}{9.98} & \pmstat{20.00}{8.16} & \pmstat{67.50}{3.54} & \pmstat{19.33}{8.22} & \pmstat{28.00}{1.63} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $1 \times 384$ & \xmark & \pmstat{28.67}{0.94} & \pmstat{76.00}{4.90} & \pmstat{34.67}{1.89} & \pmstat{79.33}{2.49} & \pmstat{18.00}{1.63} & \pmstat{46.00}{3.27} & \pmstat{2.00}{1.63} & \pmstat{11.33}{3.40} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $1 \times 384$ & \cmark & \pmstat{42.00}{3.27} & \pmstat{56.67}{4.11} & \pmstat{38.67}{3.40} & \pmstat{96.00}{0.00} & \pmstat{22.67}{5.73} & \pmstat{78.00}{2.83} & \pmstat{5.33}{3.40} & \pmstat{14.67}{0.94} \\
    \rowcolor{Blue!5} ResNet (from scratch) & $1 \times 384$ & \xmark & \pmstat{4.67}{3.40} & \pmstat{10.00}{1.63} & \pmstat{82.00}{8.49} & \pmstat{96.00}{2.83} & \pmstat{66.00}{2.83} & \pmstat{91.33}{0.94} & \pmstat{2.00}{2.83} & \pmstat{29.33}{3.40} \\
    \rowcolor{Blue!5} ResNet (from scratch) & $1 \times 384$ & \cmark & \pmstat{84.00}{7.12} & \pmstat{\textbf{100.00}}{0.00} & \pmstat{52.00}{6.53} & \pmstat{86.67}{0.94} & \pmstat{54.00}{7.12} & \pmstat{98.00}{0.00} & \pmstat{19.33}{3.40} & \pmstat{48.67}{4.99} \\
    DINOv2 (patch) & $14 \times 14 \times 384$ & \xmark & \pmstat{52.67}{5.73} & \pmstat{76.67}{6.18} & \pmstat{35.33}{4.11} & \pmstat{80.67}{6.18} & \pmstat{40.83}{10.07} & \pmstat{76.67}{5.14} & \pmstat{56.00}{4.32} & \pmstat{66.00}{4.90} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $14 \times 14 \times 8$ & \xmark & \pmstat{80.00}{7.12} & \pmstat{90.67}{3.77} & \pmstat{44.00}{7.12} & \pmstat{81.33}{6.80} & \pmstat{72.00}{4.32} & \pmstat{96.67}{0.94} & \pmstat{70.00}{1.63} & \pmstat{78.67}{0.94} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $14 \times 14 \times 8$ & \cmark & \pmstat{\textbf{90.67}}{0.94} & \pmstat{\textbf{100.00}}{0.00} & \pmstat{\textbf{94.00}}{1.63} & \pmstat{\textbf{100.00}}{0.00} & \pmstat{\textbf{82.67}}{3.77} & \pmstat{98.67}{0.94} & \pmstat{\textbf{77.33}}{6.18} & \pmstat{85.33}{4.99} \\
    \rowcolor{Blue!5} ResNet (from scratch) & $14 \times 14 \times 8$ & \xmark & \pmstat{1.33}{1.89} & \pmstat{6.67}{1.89} & \pmstat{14.67}{4.99} & \pmstat{66.00}{9.09} & \pmstat{18.67}{4.11} & \pmstat{57.33}{4.71} & \pmstat{71.33}{7.36} & \pmstat{70.67}{10.50} \\
    \rowcolor{Blue!5} ResNet (from scratch) & $14 \times 14 \times 8$ & \cmark & \pmstat{84.67}{2.49} & \pmstat{\textbf{100.00}}{0.00} & \pmstat{64.67}{8.38} & \pmstat{98.67}{1.89} & \pmstat{80.67}{0.94} & \pmstat{\textbf{99.33}}{0.94} & \pmstat{70.67}{0.94} & \pmstat{\textbf{91.33}}{2.49} \\
    \bottomrule
  \end{tabular}%
  }

\end{table*}
```
```{=latex}
\begin{figure*}[t]\center
    \begin{subfigure}[t]{0.24\textwidth}

        \includegraphics[width=\textwidth]{images/wall_dim8_hw14_mpc.pdf}

        \caption{Wall}
        \label{fig:mpc_wall}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.24\textwidth}

        \includegraphics[width=\textwidth]{images/umaze_dim8_hw14_mpc.pdf}

        \caption{PointMaze-UMaze}
        \label{fig:mpc_umaze}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.24\textwidth}

        \includegraphics[width=\textwidth]{images/medium_dim8_hw14_mpc.pdf}

        \caption{PointMaze-Medium}
        \label{fig:mpc_medium}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.24\textwidth}

        \includegraphics[width=\textwidth]{images/pusht_dim8_hw14_mpc.pdf}

        \caption{PushT}
        \label{fig:mpc_pusht}
    \end{subfigure}

\caption{Success Rate over MPC Steps. The dashed black lines represent DINO-WM with frozen DINOv2 patch features. The solid lines represent frozen DINOv2 patch features with a trainable channel projector (with resulting features $z \in \mathbb{R}^{14\times14\times 8}$) with different strengths of straightening. Our model reaches 100\% success rates very quickly as shown in~\Cref{fig:mpc_wall,fig:mpc_umaze}.
}

\label{img:mpc}
\end{figure*}
```
## Planning

We show that straightening can significantly improve the planning success rates across models and environments.

#### Setup.

The start states and goals are sampled from test trajectories to guarantee that goals can be reached within 25 steps. We follow DINO-WM [@dinowm] in using a frameskip of five, so we only need to roll out the world model for $H=5$ times. During test time, an action sequence is optimized using gradient descent through the learned dynamics model ($f_\theta$) to minimize a goal cost. For PushT, we assume we have both target images and proprioceptions; for other environments, we only use target images to increase the task difficulty.

We evaluate performance in both open-loop and closed-loop settings. Open-loop planning optimizes a length-$H$ action sequence using the MSE between the terminal embedding and the target embedding as the planning cost. Closed-loop MPC replans at every step: it optimizes a length-$H$ action sequence, executes only the first action, and then replans, using a weighted objective that encourages the predicted trajectory to approach the target (see `\Cref{appendix:planning}`{=latex}). For PushT specifically, we use only the terminal loss within horizon $H$ because regime-switching dynamics make intermediate-state loss misleading, so we apply the weighted intermediate loss only beyond $H$ where it is more stable.

#### Results.

As shown in `\Cref{tab:main_results}`{=latex}, we observe a significant improvement across all models and environments. When training the projectors or encoders, we observe an improvement in performance even without the straightening regularization. We attribute this improvement to the *implicit straightening* during training as discussed in `\Cref{sec:analysis_emb}`{=latex}. For ResNet with spatial features, we observe abnormally low success rates for Wall, PointMaze-UMaze and PointMaze-Medium, which could be explained by the extremely high curvature in `\Cref{img:cos}`{=latex}, suggesting a degradation of features. We also notice that the implicit straightening is the weakest for the UMaze when using the projector, which also results in the lowest improvement in planning.

Applying explicit straightening further strengthens the straightness in the embedding space, resulting in more than 10% boost in open-loop and MPC success rates for almost all setups. For example, UMaze's open-loop success rate is improved from 44% to 94% with the projector, and 14.67% to 64.67% when training a ResNet from scratch. Note that we use weighted loss on intermediate states which enables reaching the target before consuming the full horizon $H=5$. It is impressive that our model reaches 100% success with MPC on Wall and UMaze within only a few steps (`\Cref{img:mpc}`{=latex}), suggesting it discovers more direct trajectories than the randomly generated test trajectories. The PushT success increases more slowly because we apply only the terminal loss within the horizon $H=5$, yet straightening still yields substantial final gains. We also compare with other widely used temporal regularization, namely smoothness and temporal contrastiveness, in `\Cref{app:smooth_tc}`{=latex}, but find temporal straightening significantly more effective.

```{=latex}
\begin{figure*}[t]

\noindent
\makebox[\linewidth][c]{%

    \makebox[\linewidth][r]{\rotatebox[origin=c]{90}{\footnotesize (a) Success}}
%
%

    \includegraphics[width=\linewidth]{images/pusht_long_good.pdf}
%
}



\makebox[\linewidth][c]{%

    \makebox[\linewidth][r]{\rotatebox[origin=c]{90}{\footnotesize (b) Failure}}
%
%

    \includegraphics[width=\linewidth]{images/pusht_long_bad.pdf}
%
}
\caption{Examples of Long-Horizon Open-Loop GD Planning on PushT. For each example, the top row shows simulator-rendered images and the bottom row shows decoded images, with the last column being the target. The failure example shows a case where the long-horizon imagined rollout does not match the real dynamics.}
\label{img:long}

\end{figure*}
```
#### Comparison of CEM and GD.

We compare the open-loop performance of gradient descent (GD) and the cross-entropy method (CEM) in `\Cref{app:ablations_subplanner}`{=latex}. Straightening regularization consistently improves the success rate of both planners. For example, on Wall and PushT, it improves the projector baseline by roughly 10% for both GD and CEM. Overall, CEM achieves higher success rates but requires substantially longer planning time than GD. With straightening, GD achieves a better success--latency trade-off.

#### Effect of feature dimensions.

We find that preserving spatial structure generally matters more than retaining channels. When we keep all patch tokens, we can aggressively reduce the channel dimension of DINOv2 features from 384 down to 8 without degrading performance. Increasing the channel dimension to $d\in\{32, 128\}$ does not improve performance and can even lead to a drop for some environments (`\Cref{app:ablations_dim}`{=latex}), which is not surprising as lower dimensions can simplify both dynamics prediction and downstream optimization. In contrast, collapsing patch features into a single global vector makes precise dynamics prediction harder. The predictor produces less accurate rollouts, which in turn reduces planning success. Notably, training a ResNet from scratch produces significantly better global features than training a global projector on frozen DINO patch features.

#### Long horizon.

To further stress test our method, we also evaluate a longer-horizon setting where the target is 50 steps away. We leave out UMaze and Wall, because in those environments, a target picked via random 50-step rollouts can end up surprisingly close in terms of shortest-path distance, which does not reflect true long-horizon difficulty. We summarize the results in `\Cref{tab:long_horizon}`{=latex} and show success and failure examples in `\Cref{img:long}`{=latex}. As expected, success rates drop substantially compared to the short-horizon setting, but our method consistently outperforms the baseline across all settings. More broadly, long-horizon rollouts remain a well-known challenge for latent planning where prediction errors compound over steps and lead to substantial trajectory drift. This is visible in failure cases where decoded rollouts become blurry or misaligned with the simulator.

Motivated by `\Cref{img:heatmap}`{=latex}, where the aggregation head produces a smoother long-range distance signal than spatial features alone, we add a global goal cost for long-horizon planning. Specifically, we keep the spatial goal cost and add a goal cost computed in the aggregated feature space: $\mathcal{L}_{\mathrm{plan}}=\mathcal{L}_{\mathrm{spatial}}+0.1\mathcal{L}_{\mathrm{agg}}$. Here, $\mathcal{L}_{\mathrm{spatial}}$ measures squared goal distance over spatial features, while $\mathcal{L}_{\mathrm{agg}}$ measures squared goal distance after applying the aggregation head. As shown in the last two rows of `\Cref{tab:long_horizon}`{=latex}, this combined cost improves over using the spatial cost alone across all models under MPC. These results suggest that long-horizon planning may benefit from objectives that combine fine-grained local costs with global distance geometry.

#### Teleported-PointMaze.

Pretrained visual embeddings primarily reflect visual similarity, whereas our straightening objective is designed to align the latent space with temporal dynamics. To test whether straightening truly captures dynamics rather than exploiting appearance cues, we introduce *Teleported-PointMaze* with modified transitions: touching the right wall instantly teleports the agent to the left side (see `\Cref{appendix:teleportation}`{=latex}). This creates states that are far in the pixel space but have small temporal distance. We visualize a representative success case in `\Cref{img:teleport_str_comp}`{=latex}, where the straightened model plans to reach the target by leveraging teleportation.

#### Limitations and future directions.

Our current formulation focuses on continuous goal-conditioned latent planning with a symmetric Euclidean goal cost, which may be suboptimal under asymmetric or irreversible dynamics. While our curvature regularizer straightens observed latent transitions and does not itself assume reversibility, such settings may require directional planning costs such as quasimetrics. Furthermore, the gains from using an aggregation head for long-horizon planning suggest that regularization and planning objectives do not necessarily operate in the prediction latent space: the world model can learn dynamics in one space, while the planner optimizes a task- and geometry-aware objective in a projected space.

```{=latex}
\setlength{\tabcolsep}{5pt}
```
```{=latex}
\renewcommand{\arraystretch}{1.08}
```
```{=latex}
\resizebox{\linewidth}{!}{%
  \begin{tabular}{lc|cc|cc}
    \toprule
    Model & $\mathcal{L}_{curv}$ &
    PushT &
    PointMaze -- Medium \\
    \cmidrule(lr){3-4}\cmidrule(lr){5-6}
     & & Open-loop & MPC & Open-loop & MPC \\
    \midrule
    DINO-WM & -- & 3.33 $\pm$ {\tiny 2.36} & 27.33 $\pm$ {\tiny 6.66} & 35.00 $\pm$ {\tiny 2.35} & 65.33 $\pm$ {\tiny 3.13} \\
    + Proj & \xmark & 6.67 $\pm$ {\tiny 3.77} & 26.67 $\pm$ {\tiny 9.98} & 60.00 $\pm$ {\tiny 3.27} & 72.00 $\pm$ {\tiny 0.00} \\
    + Proj & \cmark & 13.33 $\pm$ {\tiny 3.77} & 24.00 $\pm$ {\tiny 6.53} & 68.00 $\pm$ {\tiny 8.64} & 88.00 $\pm$ {\tiny 3.27} \\
    ResNet & \xmark & 13.33 $\pm$ {\tiny 3.77} & 29.33 $\pm$ {\tiny 9.43} & 14.67 $\pm$ {\tiny 6.80} & 48.00 $\pm$ {\tiny 9.80} \\
    ResNet & \cmark & 10.67 $\pm$ {\tiny 4.99} & 33.33 $\pm$ {\tiny 4.99} & 76.00 $\pm$ {\tiny 6.53} & 98.67 $\pm$ {\tiny 1.89} \\
    \midrule
    \emph{Combined planning cost:} $\mathcal{L}_{\mathrm{plan}}=\mathcal{L}_{\mathrm{spatial}}+0.1\mathcal{L}_{\mathrm{agg}}$ \\
    \midrule
     + Proj \, & \cmark & 20.00 $\pm$ {\tiny 0.00} & 33.33 $\pm$ {\tiny 4.16} & 66.67 $\pm$ {\tiny 7.57} & 92.00 $\pm$ {\tiny 5.29} \\
    ResNet   & \cmark & 13.33 $\pm$ {\tiny 4.62} & 36.00 $\pm$ {\tiny 5.29} & 68.67 $\pm$ {\tiny 4.16} & 98.67 $\pm$ {\tiny 1.15} \\

    \bottomrule
  \end{tabular}%
  }
```
# Conclusion

In this work, we show that temporal straightening yields an embedding space that effectively facilitates latent planning. In this straightened representation space, the Euclidean distance provides a more reliable proxy for the geodesic distance and gradient-based planning is better conditioned. Across a range of 2D goal-reaching tasks, this leads to significant and consistent gains over baselines. More broadly, our findings highlight that representation geometry plays an important role in latent planning and show that straightening latent trajectories is a simple yet effective way to improve it. We believe this opens a promising path toward more efficient latent planning in more challenging environments.

# Impact Statement {#impact-statement .unnumbered}

This paper presents work whose goal is to advance the field of Machine Learning. World models with improved planning capabilities could have both beneficial applications (e.g., robotics, autonomous systems, scientific discovery) and potential risks if deployed without adequate safety measures. We encourage future work to consider safety implications when deploying such systems in real-world settings.

# Acknowledgments {#acknowledgments .unnumbered}

We thank Yilun Kuang and Daohan Lu for helpful discussions. This work was supported in part by AFOSR under grant FA95502310139, NSF Award 1922658, Visko AI, a Google TPU Award, the NYU-KAIST Award A25-0081-002, and the Institute of Information & Communications Technology Planning Evaluation (IITP) under grant RS-2024-00469482, funded by the Ministry of Science and ICT (MSIT) of the Republic of Korea in connection with the Global AI Frontier Lab International Collaborative Research. The compute is supported by the NYU High Performance Computing resources, services, and staff expertise.

```{=latex}
\bibliographystyle{icml2026}
```
```{=latex}
\newpage
```
```{=latex}
\appendix
```
```{=latex}
\onecolumn
```
```{=latex}
\appendix
```
# Appendix {#appendix .unnumbered}

# Data and Environments {#app:datasets}

## Wall

This is a 2D navigation environment introduced by @dinowm and @sobal2025learning. The environment consists of two rooms separated by a wall with a single narrow door. To move between rooms, the agent must pass through this door. The task is to navigate from a start position to a target position, given start and target images. The action space consists of 2D vectors representing displacements in x and y axes. For training, we follow the approach of @dinowm to generate a dataset of 1,920 trajectories, each 50 time steps long. We train for 20 epochs.

## PointMaze (UMaze and Medium-Maze) {#appendix:pointmze}

This is a 2D navigation environment based on the MuJoCo physics engine [@fu2020d4rl]. We experiment on the \`\`UMaze" and \`\`Medium-Maze" here and plan to test other maze setups in future work. The task is to navigate from a start position to a target position, given start and target images. Unlike the previous \`\`Wall" environment, this dynamics is governed by realistic physical properties such as velocity, acceleration, and inertia. The action space consists of forces applied along the x and y axes. For training, we follow @dinowm to generate a dataset of 2,000 trajectories for UMaze and 4,000 for Medium-Maze, each 100 time steps long. We train for 20 epochs.

## PushT {#appendix:pointmze}

This is a challenging, contact-rich environment introduced by @chi2024pusht. PushT features a pusher agent interacting with a T-shaped block. Starting from a random initial state, the agent must drive both the pusher and the T-block to a known feasible target configuration matching their target poses. The fixed green T is not the T-block's target and is only a visual reference marker. We use training data from @dinowm, which contains 18500 trajectories with lengths of 100-300. We train for 2 epochs.

# Experiments {#app:experiments}

## Model Predictive Control (MPC) {#appendix:planning}

We outline the MPC algorithm below. Unlike DINO-WM [@dinowm] that uses Cross-Entropy Method (CEM) as the subplanner, we use gradient descent for the major experiments instead to accelerate planning.

a\) **Encode States:** Given the current observation $o_0$ and the goal observation $o_g$ (both RGB images), we first encode them into their latent state representations using our trained encoder $\mathcal{E}^s$ (either a pre-trained DINOv2 encoder plus a projector, or a ResNet from scratch): $$z_0 = \mathcal{E}^s(o_0), \quad z_g = \mathcal{E}^s(o_g).$$

b\) **Initialize Actions:** An initial action sequence for the planning horizon $T$ is sampled from Gaussian distribution, $\{a_0, a_1, \dots, a_{T-1}\}$.

c\) **Define Objective:** The planning objective is to minimize the mean squared error (MSE) between the predicted final latent state $\hat{z}_T$ and the goal state $z_g$: $$L = \|\hat{z}_T - z_g\|_2^2$$ where the latent trajectory is predicted by recursively applying the world model: $\hat{z}_t = f_{\theta}(\hat{z}_{t-1}, a_{t-1})$.

d\) **Optimize via Gradient Descent:** Update actions iteratively using gradients of the cost with respect to the actions: $$a_t \leftarrow a_t - \eta \frac{\partial L}{\partial a_t}, \quad \text{for } t = 0, \dots, T-1,$$ where $\eta$ is the learning rate. Repeat until reaching the predefined number of iterations.

e\) **Execute Action:** After the optimization loop is complete, the first $k$ actions from the optimized action sequence are executed in the environment.

f\) **Re-plan:** The process is repeated from step (a) at the next environment timestep, using the new observation $o_1$.

```{=latex}
\clearpage
```
## Hyperparameters

::: {#tab:plan_hyperparams}
  **Name**                 **Value**
  ------------------------ -----------
  Projector/ResNet lr      1e-5[^2]
  Predictor lr             5e-4
  Action/Prop encoder lr   5e-4
  Batch size               32
  History frames           3
  Frameskip                5

  : Planning Hyperparameters.
:::

```{=latex}
\hfill
```
::: {#tab:plan_hyperparams}
  **Name**                **Value**
  ----------------------- -----------
  Subplanner horizon      25
  \# Executed actions     25[^3]
  Optimizer               Adam
  Action Initialization   Zero
  Learning rate           0.1
  #opt steps              100

  : Planning Hyperparameters.
:::

## Planning: GD vs. CEM {#app:ablations_subplanner}

```{=latex}
\begin{wrapfigure}{r}{0.45\textwidth}


\includegraphics[width=\linewidth]{images/speed_vs_success.pdf}
\caption{Success rate versus wall-clock planning time for open-loop GD and CEM planning after straightening.}

\label{img:time}

\end{wrapfigure}
```
We compare the open-loop success rate using GD and CEM planners. For one single plan, GD optimizes the action sequence by backpropagating through the learned rollout model. With $N$ optimization steps, this requires $N$ forward rollouts and $N$ backward passes. In contrast, CEM iteratively samples $M$ candidate action sequences, rolls each candidate out with the learned predictor, refits the sampling distribution to the top-performing candidates, and repeats this procedure for $K$ iterations. Thus, CEM requires $MK$ forward rollouts, with $M$ typically large for competitive performance.

In our experiments, we find that CEM requires at least $200$ samples and $10$ iterations to achieve strong performance, making it roughly $10\times$ slower than GD in wall-clock planning time. We report wall-clock time for open-loop planning over 50 test trajectories on a single L40S GPU in `\Cref{img:time}`{=latex}.

As shown in `\Cref{tab:open_loop_gd_cem}`{=latex}, straightening consistently improves **both** GD and CEM across environments and model architectures. Consistent with prior work [@dinowm], CEM often obtains higher absolute success rates than GD, but at substantially higher computational cost. Importantly, straightening largely reduces the performance gap between GD and CEM, suggesting that the regularizer improves the latent optimization landscape and enables simple gradient-based planning to achieve a better success--latency trade-off.

```{=latex}
\begin{table*}[h]

  \caption{Goal-reaching Success Rate of 50 Test Trajectories (\%) in open-loop planning. We compare GD and CEM planners. Values are mean $\pm$ std over three data seeds. The best value is \textbf{bold}.}
  \label{tab:open_loop_gd_cem}
  \setlength{\tabcolsep}{5pt}
  \renewcommand{\arraystretch}{1.08}
  \resizebox{\linewidth}{!}{
  \begin{tabular}{lcc@{}|@{}cc@{}cc@{}cc@{}cc}
    \toprule
    Method & Config &
    Wall &
    PointMaze -- UMaze &
    PointMaze -- Medium &
    PushT \\
    \cmidrule(lr){2-3}\cmidrule(lr){4-5}\cmidrule(lr){6-7}\cmidrule(lr){8-9}\cmidrule(lr){10-11}
     & dim & $\mathcal{L}_{curv}$ & GD & CEM & GD & CEM & GD & CEM & GD & CEM \\
    \midrule
    DINOv2 (patch) & $14 \times 14 \times 384$ & \xmark & 73.33 $\pm$ {\tiny 3.40} & 87.33 $\pm$ {\tiny 4.99} & 63.33 $\pm$ {\tiny 8.22} & 88.00 $\pm$ {\tiny 1.63} & 70.00 $\pm$ {\tiny 4.08} & 88.00 $\pm$ {\tiny 1.63} & 62.67 $\pm$ {\tiny 4.11} & 71.33 $\pm$ {\tiny 7.72} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $14 \times 14 \times 8$ & \xmark & 80.00 $\pm$ {\tiny 7.12} & 92.00 $\pm$ {\tiny 0.00} & 44.00 $\pm$ {\tiny 7.12} & 75.33 $\pm$ {\tiny 4.99} & 72.00 $\pm$ {\tiny 4.32} & \textbf{92.67} $\pm$ {\tiny 4.71} & 70.00 $\pm$ {\tiny 1.63} & 71.33 $\pm$ {\tiny 6.18} \\
    \rowcolor{Blue!5} DINOv2 (patch) + proj & $14 \times 14 \times 8$ & \cmark & \textbf{90.67} $\pm$ {\tiny 0.94} & \textbf{100.00} $\pm$ {\tiny 0.00} & \textbf{94.00} $\pm$ {\tiny 1.63} & \textbf{94.00} $\pm$ {\tiny 1.63} & \textbf{82.67} $\pm$ {\tiny 3.77} & 86.67 $\pm$ {\tiny 1.89} & \textbf{77.33} $\pm$ {\tiny 6.18} & \textbf{80.00} $\pm$ {\tiny 4.32} \\
    \rowcolor{Blue!5} ResNet & $14 \times 14 \times 8$ & \xmark & 1.33 $\pm$ {\tiny 1.89} & 1.33 $\pm$ {\tiny 0.94} & 14.67 $\pm$ {\tiny 4.99} & 20.67 $\pm$ {\tiny 0.94} & 18.67 $\pm$ {\tiny 4.11} & 24.00 $\pm$ {\tiny 4.32} & 71.33 $\pm$ {\tiny 7.36} & 56.00 $\pm$ {\tiny 0.00} \\
    \rowcolor{Blue!5} ResNet & $14 \times 14 \times 8$ & \cmark & 84.67 $\pm$ {\tiny 2.49} & 90.00 $\pm$ {\tiny 5.89} & 64.67 $\pm$ {\tiny 8.38} & 83.33 $\pm$ {\tiny 2.49} & 80.67 $\pm$ {\tiny 0.94} & 89.33 $\pm$ {\tiny 6.18} & 70.67 $\pm$ {\tiny 0.94} & 72.67 $\pm$ {\tiny 6.60} \\
    \bottomrule
  \end{tabular}
  }
\end{table*}
```
```{=latex}
\newpage
```
## Effect of Feature Dimensions {#app:ablations_dim}

In order to improve efficiency and efficacy, we ablate the output dimensions of the encoders. Here, we test on the \`\`frozen DINOv2 + spatial projector" setup and preserve the spatial dimensions of the DINOv2 patch features $m_v=196$ but decreasing channels from 384 to $d_v \in \{2,8,32,128\}$. For all experiments, we use $lr=1e-6$ for the encoder. If with straightening, we apply straightening on the aggregation head as described in `\Cref{app:straightening}`{=latex} with a straightening strength $\lambda=0.1$.

We report the open-loop planning success rate of 50 test samples over three data sampling seeds in `\Cref{img:ablation_dim}`{=latex}. Very small dimensions (e.g., $d_v=2$) result in poor performance, indicating insufficient capacity to preserve planning-relevant information. Increasing to a moderate dimension ($d_v=\{8,32\}$) yields the best results, while too large dimensions ($d_v = 128$) consistently reduce success rates. This suggests that overly high-dimensional latents can hinder gradient-based planning.

<figure id="img:ablation_dim">
<figure id="fig:wall">
<img src="images/wall_both_projchannel_hw14_dim_curve.png" />
<figcaption>Wall</figcaption>
</figure>
<figure id="fig:umaze">
<img src="images/umaze_both_projchannel_hw14_dim_curve.png" />
<figcaption>PointMaze-UMaze</figcaption>
</figure>
<figure id="fig:medium">
<img src="images/medium_both_projchannel_hw14_dim_curve.png" />
<figcaption>PointMaze-Medium</figcaption>
</figure>
<figure id="fig:mpc_pusht">
<img src="images/pusht_both_projchannel_hw14_dim_curve.png" />
<figcaption>PushT</figcaption>
</figure>
<figcaption>Comparison of Different Dimensions. The line plots show the success rate changes with increasing channels. Too small dimensions (e.g. <span class="math inline"><em>d</em><sub><em>v</em></sub> = 2</span>) are unable to encode sufficient planning-relevant information, while unnecessarily high dimensions (e.g. <span class="math inline"><em>d</em><sub><em>v</em></sub> = 128</span>) hinder the planning performance. </figcaption>
</figure>

## Comparison to Smoothness and Temporal Contrastive Objectives {#app:smooth_tc}

```{=latex}
\begin{wrapfigure}{r}{0.3\textwidth}


\includegraphics[width=\linewidth]{images/smooth_tc_ablation.pdf}
\caption{Comparison with Other Regularizations.}
\label{img:smmoth_tc}
\end{wrapfigure}
```
We compare with two common temporal regularization objectives:

-   **Smoothness.** This objective penalizes large temporal jumps in visual embeddings: $$\mathcal{L}_{\mathrm{smooth}}=\mathbb{E}_t\left[\|z_{t+1}-z_t\|_2^2\right].$$ However, an overly strong smoothness penalty can lead to degenerate solutions where embeddings of different states collapse to similar values.

-   **Time contrastiveness.** This objective treats frames within a temporal window of size $k$ as positives and other frames in the same trajectory as negatives, encouraging temporally nearby embeddings to be similar and temporally distant embeddings to be different: $$\mathcal{L}_{\mathrm{tc}}=\mathrm{InfoNCE}(z_{\mathrm{pos}},z_{\mathrm{neg}}).$$ However, when training trajectories are suboptimal, temporal distance may not reflect geodesic distance: states that are geodesically close can be temporally far apart, and this objective may incorrectly separate them.

We test on the \`\`frozen DINOv2 + spatial projector" setup and report open-loop GD planning success rate on PushT in `\Cref{img:smmoth_tc}`{=latex}, evaluated on 50 test samples over three data seeds. Overall, we do not observe improvements from adding smoothness or temporal contrastive objectives. Larger weights generally hurt performance, while smaller weights are less harmful but still do not match the gains from straightening. These objectives may still be useful in other settings, and our method is complementary to them: the curvature regularization loss can be combined with any other losses.

```{=latex}
\clearpage
```
## Cosine Similarity Variants for Spatial Features {#app:straightening}

```{=latex}
\begin{wrapfigure}{r}{0.45\textwidth}


\includegraphics[width=\linewidth]{images/train_agg.pdf}
\caption{\colorbox{orange!20}{Aggregation Head} for Straightening during Training. The prediction loss is applied to spatial features, while the curvature loss is applied to the aggregated features.}
\label{img:agg}

\end{wrapfigure}
```
For spatial visual features $z_t^v \in \mathbb{R}^{m_v \times d_v}$ ($m_v>1$), we compute straightness from approximate latent velocities $v_t := z_{t+1}^v - z_t^v \in \mathbb{R}^{m_v \times d_v}.$ Let $v_{t,i}\in\mathbb{R}^{d_v}$ denote the $i$-th patch vector and $\cos(u,w)=\frac{u^\top w}{\|u\|_2\|w\|_2}$. We ablate four choices of $\mathcal{C}_t$:

-   **\[patch\]** We treat each patch independently, then average: $$\mathcal{C}_t=\frac{1}{m_v}\sum_{i=1}^{m_v}\cos\!\left(v_{t,i},\,v_{t+1,i}\right).$$

-   **\[mean\]** We average patches to one vector, then cosine: $$\bar v_t=\frac{1}{m_v}\sum_{i=1}^{m_v} v_{t,i},
    \qquad
    \mathcal{C}_t=\cos(\bar v_t,\bar v_{t+1}).$$

```{=html}
<!-- -->
```
-   **\[flatten\]** We flatten the spatial features and compute a single cosine over all dimensions: $$\mathcal{C}_t=\cos\!\left(\mathrm{vec}(v_t),\,\mathrm{vec}(v_{t+1})\right),$$ where $\mathrm{vec}(\cdot):\mathbb{R}^{m_v\times d_v}\to\mathbb{R}^{m_v d_v}$.

-   **\[agg\]** We learn an aggregation head to aggregate features to a single global feature before cosine (`\Cref{img:agg}`{=latex}): $$\mathcal{C}_t=\cos\!\left(h_\phi(v_t),\,h_\phi(v_{t+1})\right),$$ with an aggregation head $h_\phi:\mathbb{R}^{m_v\times d_v}\to\mathbb{R}^{d_h}$. Concretely, we use an MLP with an output dimension of 128 as $h_\phi$ in all experiments.

We test these variants on the \`\`frozen DINOv2 + spatial projector" setup and report the open-loop planning success rate of 50 test samples over three data sampling seeds in `\Cref{img:ablation_str}`{=latex}. The projector projects pretrained DINOv2 patch features $e_t \in \mathbb{R}^{196\times384}$ to $z_t \in \mathbb{R}^{196\times8}$. For the straightening strength coefficient, we use $\lambda=0.1$ for agg and $\lambda=0.01$ for the rest, as these values yield the best performance. We find that using a learnable aggregation head performs best. This is not surprising as straightening should act on the *global* trajectory representations, whereas spatial tokens mainly capture local, patch-level variations that are only loosely aligned across time due to object motion and occlusion.

<figure id="img:ablation_str">
<figure id="fig:wall">
<img src="images/wall_dim8_hw14_straightening_bar.png" />
<figcaption>Wall</figcaption>
</figure>
<figure id="fig:umaze">
<img src="images/umaze_dim8_hw14_straightening_bar.png" />
<figcaption>PointMaze-UMaze</figcaption>
</figure>
<figure id="fig:medium">
<img src="images/medium_dim8_hw14_straightening_bar.png" />
<figcaption>PointMaze-Medium</figcaption>
</figure>
<figure id="fig:mpc_pusht">
<img src="images/pusht_dim8_hw14_straightening_bar.png" />
<figcaption>PushT</figcaption>
</figure>
<figcaption>Comparison of Different Straightening Strategies. The bar charts show the planning success rates. While all cosine similarity variants lead to better performance than no straightening, adding a learnable aggregation head gives the best performance. </figcaption>
</figure>

```{=latex}
\clearpage
```
# Theoretical Analysis

## Setup and notation

We optimize an action sequence $\mathbf a=(a_0,\dots,a_{K-1})\in\mathbb{R}^{K\times d_a}$ over horizon $K$ to minimize the terminal MSE $$\label{eq:app_obj}
\mathcal L(\mathbf a)=\|z_K-z_g\|_2^2,\qquad z_K=\Phi(\mathbf a),$$ where $\Phi$ denotes unrolling the latent dynamics from a fixed initial state $z_0$.

::: assumption
**Assumption 6** (Linear latent dynamics). `\label{as:app_linear}`{=latex} We assume linear latent dynamics $$\label{eq:app_linear}
z_{t+1}=Az_t+Ba_t,\qquad A\in\mathbb{R}^{d\times d},\; B\in\mathbb{R}^{d\times d_a}.$$
:::

::: definition
**Definition 7** (Effective condition number). `\label{def:app_keff}`{=latex} For a PSD matrix $H\succeq 0$ with a nontrivial nullspace, define $$\kappa_{\mathrm{eff}}(H):=\frac{\sigma_{\max}(H)}{\sigma_{\min}^+(H)},$$ where $\sigma_{\min}^+(H)$ is the smallest nonzero singular value.
:::

::: definition
**Definition 8** ($\varepsilon$-straight transition). `\label{def:app_eps}`{=latex} In the linear model `\eqref{eq:app_linear}`{=latex}, define $$\varepsilon:=\|A-I\|_2.$$
:::

## Conditioning of the planning Hessian {#app:theory_hessian}

Unrolling `\eqref{eq:app_linear}`{=latex} gives the affine terminal map $$\label{eq:app_unroll}
z_K
=
A^{K} z_0 + \sum_{t=0}^{K-1} A^{K-1-t} B a_t.$$ Define the rollout Jacobian $$\label{eq:app_Jphi}
J_\Phi
:=
\frac{\partial z_K}{\partial \mathbf a}
=
\bigl[\,A^{K-1}B,\quad A^{K-2}B,\quad \cdots,\quad B\,\bigr]
\in\mathbb{R}^{d\times (K d_a)}.$$ The associated finite-horizon discrete controllability Gramian is $$\label{eq:app_gram}
\mathcal W_K
:=
J_\Phi J_\Phi^\top
=
\sum_{k=0}^{K-1} A^k B B^\top (A^\top)^k\in\mathbb{R}^{d\times d},$$ a standard term in linear systems [@kailath1980linear; @sontag1998mathematical; @chen1999linear].

::: lemma
**Lemma 9** (Hessian form and Gramian equivalence). *`\label{lem:app_hess_gram}`{=latex} Under `\eqref{eq:app_obj}`{=latex}--`\eqref{eq:app_linear}`{=latex}, the planning Hessian satisfies $$\label{eq:app_hess}
H:=\nabla^2_{\mathbf a}\mathcal L(\mathbf a)=2J_\Phi^\top J_\Phi\succeq 0.$$ Moreover, the nonzero singular values of $J_\Phi^\top J_\Phi$ equal those of $J_\Phi J_\Phi^\top$, hence $$\label{eq:app_keff_equal}
\kappa_{\mathrm{eff}}(H)=\kappa(\mathcal W_K).$$*
:::

::: proof
*Proof.* $H$ is positive semi-definite by definition. Since $z_K$ is affine in $\mathbf a$ by `\eqref{eq:app_unroll}`{=latex}, $\mathcal L(\mathbf a)=\|z_K-z_g\|_2^2$ is a convex quadratic, and direct differentiation yields $H=2J_\Phi^\top J_\Phi\succeq 0$. For any matrix $M$, the nonzero eigenvalues of $M^\top M$ and $MM^\top$ coincide. Applying this with $M=J_\Phi$ gives that the nonzero eigenvalues of $H/2$ equal those of $\mathcal W_K$, which implies `\eqref{eq:app_keff_equal}`{=latex}. ◻
:::

::: theorem
**Theorem 10** (Conditioning bound). *`\label{thm:app_cond}`{=latex} Assume `\eqref{eq:app_linear}`{=latex}. Consider first the square-action case $d_a=d$ with $B$ invertible. Then $$\label{eq:app_kappaA}
\kappa_{\mathrm{eff}}(H)=\kappa(\mathcal W_K)
\;\le\;
\kappa(B)^2\,
\frac{\sum_{k=0}^{K-1}\sigma_{\max}(A)^{2k}}{\sum_{k=0}^{K-1}\sigma_{\min}(A)^{2k}}
\;\le\;
\kappa(B)^2\,\kappa(A)^{2(K-1)},$$ where $\kappa(A):=\sigma_{\max}(A)/\sigma_{\min}(A)$. If additionally $\varepsilon=\|A-I\|_2<1$, then $$\label{eq:app_eps}
\kappa_{\mathrm{eff}}(H)
\;\le\;
\kappa(B)^2\left(\frac{1+\varepsilon}{1-\varepsilon}\right)^{2(K-1)}
\;\le\;
\kappa(B)^2 e^{6\varepsilon K}
\quad(\varepsilon\le\tfrac12).$$*
:::

::: proof
*Proof.* By Lemma `\ref{lem:app_hess_gram}`{=latex}, it suffices to bound $\kappa(\mathcal W_K)$.

#### Upper bound.

For any unit vector $x\in\mathbb{R}^d$, $$x^\top \mathcal W_K x = \sum_{k=0}^{K-1}\|B^\top (A^\top)^k x\|_2^2
\le
\sum_{k=0}^{K-1}\|B\|_2^2\|A^k\|_2^2\|x\|_2^2
\le
\sigma_{\max}(B)^2\sum_{k=0}^{K-1}\sigma_{\max}(A)^{2k}.$$ Taking the maximum over $\|x\|_2=1$ yields $$\lambda_{\max}(\mathcal W_K)\le \sigma_{\max}(B)^2\sum_{k=0}^{K-1}\sigma_{\max}(A)^{2k}.$$

#### Lower bound.

Since $B$ is invertible, $\|B^\top u\|_2\ge \sigma_{\min}(B)\|u\|_2$ for all $u$. Also $\sigma_{\min}(A^k)\ge \sigma_{\min}(A)^k$. Thus for any unit $x$, $$\|B^\top(A^\top)^k x\|_2
\ge
\sigma_{\min}(B)\,\|(A^\top)^k x\|_2
\ge
\sigma_{\min}(B)\,\sigma_{\min}(A^k)\,\|x\|_2
\ge
\sigma_{\min}(B)\,\sigma_{\min}(A)^k,$$ hence $$x^\top \mathcal W_K x
\ge
\sigma_{\min}(B)^2\sum_{k=0}^{K-1}\sigma_{\min}(A)^{2k}.$$ Taking the minimum over $\|x\|_2=1$ yields $$\lambda_{\min}(\mathcal W_K)\ge \sigma_{\min}(B)^2\sum_{k=0}^{K-1}\sigma_{\min}(A)^{2k}.$$

#### Combine.

Dividing the two bounds gives the first inequality in `\eqref{eq:app_kappaA}`{=latex}. For the second, use positivity of terms: $$\frac{\sum_{k=0}^{K-1}\sigma_{\max}(A)^{2k}}{\sum_{k=0}^{K-1}\sigma_{\min}(A)^{2k}}
\le
\max_{0\le k\le K-1}\frac{\sigma_{\max}(A)^{2k}}{\sigma_{\min}(A)^{2k}}
=
\kappa(A)^{2(K-1)}.$$

#### $\varepsilon$-specialization.

If $\varepsilon=\|A-I\|_2<1$, then by Weyl's perturbation theorem, $\sigma_{\max}(A)\le 1+\varepsilon$ and $\sigma_{\min}(A)\ge 1-\varepsilon$, which implies the first inequality in `\eqref{eq:app_eps}`{=latex}. For $\varepsilon\le \tfrac12$, the standard bound $\ln\!\big(\frac{1+\varepsilon}{1-\varepsilon}\big)\le 3\varepsilon$ gives the exponential form. ◻
:::

::: remark
*Remark 11* (Low-dimensional actions $d_a<d$). `\label{rem:app_lowdim}`{=latex} If $d_a<d$, then $B$ is not invertible and $\mathcal W_K$ may be singular. All statements hold on the controllable subspace $\mathcal S_K=\mathrm{range}(\mathcal W_K)$ by replacing $\lambda_{\min}(\mathcal W_K)$ with $\lambda_{\min}^+(\mathcal W_K)$ and interpreting $\kappa(\mathcal W_K)$ as an effective condition number. In this case, additional controllability assumptions are needed to lower bound $\sigma_{\min}^+(\mathcal W_K)$.
:::

## Cosine similarity as a proxy {#app:theory_cos}

::: assumption
**Assumption 12** (Constant velocity and smooth actions). `\label{as:app_cos_const}`{=latex} Define latent velocities $v_t:=z_{t+1}-z_t$. Assume there exists a constant $c>0$ such that $$\|v_t\|_2=c\qquad \text{for all } t=0,\dots,K-1.$$ Assume action smoothness $\Delta_a:=\max_{t}\|a_{t+1}-a_t\|_2<\infty$.
:::

::: definition
**Definition 13** (Cosine similarity). `\label{def:app_cos}`{=latex} For $t=0,\dots,K-2$, define $$\mathcal C_t:=\cos(v_t,v_{t+1})
=\frac{v_t^\top v_{t+1}}{\|v_t\|_2\|v_{t+1}\|_2},
\qquad
\bar{\mathcal C}:=\frac{1}{K-1}\sum_{t=0}^{K-2}\mathcal C_t.$$
:::

::: proposition
**Proposition 14** (Cosine proxy $\Rightarrow$ small $(A-I)$ along visited directions). *`\label{prop:app_cos}`{=latex} Under linear dynamics `\eqref{eq:app_linear}`{=latex}, let $\hat v_t:=v_t/\|v_t\|_2$. Under Assumption `\ref{as:app_cos_const}`{=latex}, for each $t=0,\dots,K-2$, $$\label{eq:app_dir_point_const}
\|(A-I)\hat v_t\|_2
\;\le\;
\sqrt{2(1-\mathcal C_t)}
\;+\;
\frac{\sigma_{\max}(B)\Delta_a}{c}.$$ If $\bar{\mathcal C}\ge 1-\eta$, then $$\label{eq:app_dir_avg_const}
\frac{1}{K-1}\sum_{t=0}^{K-2}\|(A-I)\hat v_t\|_2
\;\le\;
\sqrt{2\eta}
\;+\;
\frac{\sigma_{\max}(B)\Delta_a}{c}.$$*
:::

::: proof
*Proof.* Under `\eqref{eq:app_linear}`{=latex}, $$v_{t+1}-v_t
=
(z_{t+2}-z_{t+1})-(z_{t+1}-z_t)
=
(A-I)(z_{t+1}-z_t)+B(a_{t+1}-a_t)
=
(A-I)v_t+B(a_{t+1}-a_t).$$ Thus, by the triangle inequality, $$\|(A-I)\hat v_t\|_2
=
\frac{\|(A-I)v_t\|_2}{\|v_t\|_2}
\le
\frac{\|v_{t+1}-v_t\|_2}{\|v_t\|_2}
+
\frac{\|B(a_{t+1}-a_t)\|_2}{\|v_t\|_2}
\le
\frac{\|v_{t+1}-v_t\|_2}{c}
+
\frac{\sigma_{\max}(B)\Delta_a}{c}.$$ Since $\|v_t\|_2=\|v_{t+1}\|_2=c$, $$\|v_{t+1}-v_t\|_2^2
=
\|v_{t+1}\|_2^2+\|v_t\|_2^2-2v_{t+1}^\top v_t
=
2c^2(1-\mathcal C_t),$$ hence $\|v_{t+1}-v_t\|_2/c=\sqrt{2(1-\mathcal C_t)}$, proving `\eqref{eq:app_dir_point_const}`{=latex}. Averaging and applying Jensen's inequality to the concave map $x\mapsto\sqrt{x}$ gives $$\frac{1}{K-1}\sum_{t=0}^{K-2}\sqrt{1-\mathcal C_t}
\le
\sqrt{1-\bar{\mathcal C}}
\le
\sqrt{\eta},$$ which implies `\eqref{eq:app_dir_avg_const}`{=latex}. ◻
:::

::: remark
*Remark 15* (Directional vs. spectral control). `\label{rem:app_dir_vs_spec}`{=latex} Proposition `\ref{prop:app_cos}`{=latex} bounds $(A-I)$ only along visited directions $\{\hat v_t\}$. Upgrading this to a uniform spectral bound $\varepsilon=\|A-I\|_2$ requires an additional coverage condition so that visited directions span the latent space. This is not an unreasonable assumption since training trajectories are typically collected to be diverse. Under such regimes, maximizing cosine similarity provides a meaningful proxy for making $A$ close to $I$ in spectral norm.
:::

# Related Work (Cont.)  {#sec:related_cont}

Here, we discuss the connections and differences between temporal straightening and local linearization, Koopman methods, and the broader literature on learning plannable representations. Temporal straightening targets the curvature of latent trajectories, a geometric property distinct from linear dynamics, whether local or global.

Local model-based control approximates nonlinear dynamics around a nominal trajectory using first- or second-order models, as in DDP and iLQR [@mayne1966ddp; @li2004ilqr]. Because these methods presuppose a low-dimensional state space, they have motivated representation-learning methods that map high-dimensional observations into latent spaces where local dynamics models become applicable. For example, E2C and RCE explicitly impose locally linear latent dynamics [@watter2015embedcontrollocallylinear; @banijamali2018robust], while later methods broaden this direction by learning representations or objectives that support locally linear control [@zhang2019solar; @Levine2020pcc; @shu2020pc3]. Temporal straightening differs in both goal and mechanism. We do not aim to learn locally linear dynamics, nor do we design representations for locally linear control. Instead, we jointly learn the encoder and predictor while directly regularizing the geometry of latent trajectories.

Koopman methods seek observables whose evolution is linear under a global operator [@koopman1931hamiltonian]. Recent deep models learn such observables or eigenfunctions directly from data [@lusch2018koopman; @yeung2019koopman; @takeishi2017learning; @cheng2026information]. Temporal straightening targets a different property of the learned representation. Koopman methods constrains the form of the dynamics model but do not require latent trajectories to be straight as linear systems can still produce curved or oscillatory paths. In contrast, temporal straightening allows nonlinear latent dynamics through a learned predictor, but more directly regularizes trajectory geometry by penalizing curvature along observed transitions.

More broadly, our work is related to learning plannable representations, especially methods that make planning easier by aligning latent geometry or value structure with feasible transitions. This includes approaches where planning or inference can be carried out through interpolation in latent space [@eysenbach2024inference; @kurutach2018causalinfogan; @wang2019visualplanning], as well as methods that build planning-oriented geometry from shortest-path, reachability, or asymmetric goal-reaching structure [@yang2020plan2vec; @wang2023quasimetric]. These works share the premise that representation geometry or value structure matters for planning, but with different focuses. Our contribution is to use straightening itself as a simple geometry regularization during training.

```{=latex}
\clearpage
```
# Visualizations

## Distance Heatmaps {#app:heatmap}

We plot heatmaps of the Euclidean distances in the embedding space. The yellow star represents the target, and we compute the Euclidean distance between its embedding and those of all other states in the maze. Blue indicates small values, and red indicates large values. With straightening, the latent distance accurately reflects the minimum number of steps required to reach the target. We find that spatial and global/aggregated features capture complementary distance information: spatial features preserve local geometry and yield fine-grained, locally discriminative distances, while global/aggregated features provide an informative longer-range signal.

We compare the distance heatmaps with ground-truth heatmaps constructed by dividing the mazes into discrete grids and applying the A-star algorithm. 4-neighbor connectivity means each grid cell connects only to up/down/left/right cells. 8-neighbor connectivity adds the four diagonals (up-left, up-right, down-left, down-right), so paths can cut corners diagonally and distances are usually shorter.

<figure id="img:more_heatmap_umaze">
<figure>
<img src="images/heatmap/umaze_geo8_gt.png" />
<figcaption>Ground-Truth using A-star (4 neighbors)</figcaption>
</figure>
<figure>
<img src="images/heatmap/umaze_geo4_gt.png" />
<figcaption>Ground-Truth using A-star (8 neighbors)</figcaption>
</figure>
<figure>
<img src="images/heatmap/umaze_dino_cls.png" />
<figcaption>DINOv2 CLS embedding</figcaption>
</figure>
<figure>
<img src="images/heatmap/umaze_dino_patch.png" />
<figcaption>DINOv2 patch embedding</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/umaze_False_projchannel_dim8_hw14/patch.png" />
<figcaption>DINOv2 + spatial proj [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/umaze_False_resnetplus_dim8_hw14_sgTrue_lr1e-5/patch.png" />
<figcaption>ResNet - spatial [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/umaze_False_resnet_dim384_hw1_sgTrue/global.png" />
<figcaption>ResNet - global [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/umaze_cos1e-1_resnet_dim384_hw1_sgTrue/global.png" />
<figcaption>ResNet - global [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/umaze_aggmlpcos1e-1_projchannel_dim8_hw14/20_agg.png" />
<figcaption>DINOv2 + spatial proj (agg head) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/umaze_aggmlpcos1e-1_projchannel_dim8_hw14/20_patch.png" />
<figcaption>DINOv2 + spatial proj (spatial) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/umaze_aggmlpcos1e-1_resnetplus_dim8_hw14/20_agg.png" />
<figcaption>ResNet - spatial (agg head) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/umaze_aggmlpcos1e-1_resnetplus_dim8_hw14/20_patch.png" />
<figcaption>ResNet - spatial (spatial) [straightening=True]</figcaption>
</figure>
<figcaption>Distance heatmaps of PointMaze-UMaze. </figcaption>
</figure>

<figure id="img:more_heatmap_medium">
<figure>
<img src="images/heatmap/medium_geo4_gt.png" />
<figcaption>Ground-Truth using A-star (4 neighbors)</figcaption>
</figure>
<figure>
<img src="images/heatmap/medium_geo8_gt.png" />
<figcaption>Ground-Truth using A-star (8 neighbors)</figcaption>
</figure>
<figure>
<img src="images/heatmap/medium_dino_cls.png" />
<figcaption>DINOv2 CLS embedding</figcaption>
</figure>
<figure>
<img src="images/heatmap/medium_dino_patch.png" />
<figcaption>DINOv2 patch embedding</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/medium_False_projchannel_dim8_hw14/patch.png" />
<figcaption>DINOv2 + spatial proj [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/medium_False_resnetplus_dim8_hw14_sgTrue_lr1e-5/patch.png" />
<figcaption>ResNet - spatial [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/medium_False_resnet_dim384_hw1_sgTrue/global.png" />
<figcaption>ResNet - global [straightening=False]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/medium_cos1e-1_resnet_dim384_hw1_sgTrue/global.png" />
<figcaption>ResNet - global [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/medium_aggmlpcos1e-1_projchannel_dim8_hw14/agg.png" />
<figcaption>DINOv2 + spatial proj (agg head) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/proj/medium_aggmlpcos1e-1_projchannel_dim8_hw14/patch.png" />
<figcaption>DINOv2 + spatial proj (spatial) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/medium_aggmlpcos1e-1_resnetplus_dim8_hw14/agg.png" />
<figcaption>ResNet - spatial (agg head) [straightening=True]</figcaption>
</figure>
<figure>
<img src="images/heatmap/resnet/medium_aggmlpcos1e-1_resnetplus_dim8_hw14/patch.png" />
<figcaption>ResNet - spatial (spatial) [straightening=True]</figcaption>
</figure>
<figcaption>Distance heatmaps of PointMaze-Medium. </figcaption>
</figure>

```{=latex}
\clearpage
```
## Visualization of Latent Trajectories {#app:pca}

To visualize the learned representations of the trajectories, we randomly sample trajectories with a length of 30 and plot them in 2D using PCA. Here, we use DINO CLS token embeddings and the aggregated features of our model (trained with straightening). While latent trajectories are highly curved in DINO CLS embedding space, they become significantly smoother after straightening. Additionally, we compute the MSE between the embeddings of each intermediate state and the target. The Euclidean distance is closer to the geodesic distance for straighter trajectories, and thus MSE (which is squared Euclidean distance) becomes a more useful planning cost function that can reflect the true progress towards the target. Visualizations for different environments are in `\Cref{img:more_pca_wall,img:more_pca_umaze,img:more_pca_medium,img:more_pca_pusht}`{=latex}.

```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/wall_exp2_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

    \includegraphics[width=0.95\textwidth]{images/wall_exp3_pca_mse.pdf}
    \end{subfigure}
\caption{PCA of Trajectories of Wall.}
\label{img:more_pca_wall}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/umaze_exp2_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

    \includegraphics[width=0.95\textwidth]{images/umaze_exp3_pca_mse.pdf}
    \end{subfigure}
\caption{PCA of Trajectories of PointMaze-UMaze.}
\label{img:more_pca_umaze}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp1_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp2_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

    \includegraphics[width=0.95\textwidth]{images/medium_exp3_pca_mse.pdf}
    \end{subfigure}
\caption{PCA of Trajectories of PointMaze-Medium.}
\label{img:more_pca_medium}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp1_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp2_pca_mse.pdf}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

    \includegraphics[width=0.95\textwidth]{images/pusht_exp3_pca_mse.pdf}
    \end{subfigure}
\caption{PCA of Trajectories of PushT. The overlaid figures only include five samples for readability.}
\label{img:more_pca_pusht}
\end{figure*}
```
```{=latex}
\clearpage
```
## Planning Trajectories {#app:traj}

<figure id="img:plan_traj_pusht">
<figure>
<img src="images/wall_exp1.png" style="width:95.0%" />
</figure>
<figure>
<img src="images/wall_exp2.png" style="width:95.0%" />
</figure>
<figure>
<img src="images/wall_exp3.png" style="width:95.0%" />
</figure>
<figcaption>Open-Loop Planning Trajectories of Wall. The first row is from the simulator and the second from the decoder. The last column is the goal image.</figcaption>
</figure>

```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/umaze_exp1.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/umaze_exp2.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/umaze_exp3.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/umaze_exp4.png}
    \end{subfigure}
\caption{Open-Loop Planning Trajectories of PointMaze-UMaze. The first row is from the simulator and the second from the decoder. The last column is the goal image.}
\label{img:plan_traj_umaze}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp1.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp2.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp3.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/medium_exp4.png}
    \end{subfigure}
\caption{Open-Loop Planning Trajectories of PointMaze-Medium. The first row is from the simulator and the second from the decoder. The last column is the goal image.}
\label{img:plan_traj_medium}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]

    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp1.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp2.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp3.png}

    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=0.95\textwidth]{images/pusht_exp4.png}
    \end{subfigure}
\caption{Open-Loop Planning Trajectories of PushT. The first row is from the simulator and the second from the decoder.}
\label{img:plan_traj_pusht}
\end{figure*}
```
```{=latex}
\clearpage
```
# Teleported-PointMaze {#appendix:teleportation}

This is a novel 2D navigation environment adapted from PointMaze. The core modification is a one-way teleportation dynamic. While the top, bottom, and left boundaries of the maze function as standard solid obstacles, a predefined region near the right wall acts as a teleportation trigger. If an agent's state transition at time $t$ results in a new x-position $x_{t+1}$ that crosses this threshold (i.e., $x_{t+1} > x_{\text{right-border}}$), an instantaneous state intervention occurs, modifying the agent's state as follows:

1.  Position (x): The agent's x-position is reset to the left side of the maze: $x_{t+1} \leftarrow x_{\text{left-border}}$.

2.  Position (y): The agent's y-position $y_{t+1}$ is preserved.

3.  Velocity (x): The agent's x-axis velocity $v_x$ is reset to its absolute value: $v_{x, t+1} \leftarrow |v_{x, t}|$.

```{=latex}
\begin{figure*}[h]\center
\includegraphics[width=\textwidth]{images/example_teleport.pdf}
\caption{Teleported-PointMaze. Note that the teleportation happens inside the red box.
}
\label{img:teleport}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]\center
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/umaze_dino_patch.pdf}
        \caption{DINOv2 patch embedding}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/heatmap_teleport_strcos.pdf}
        \caption{Trained projector with straightening}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/heatmap_teleport_strFalse.pdf}
        \caption{Trained projector without straightening}
    \end{subfigure}
    \hfill
    \begin{subfigure}[t]{0.46\textwidth}

        \includegraphics[width=\textwidth]{images/heatmap/heatmap_teleport_gt.pdf}
        \caption{Ground-Truth using A-star}
    \end{subfigure}
\caption{Distance heatmaps of Teleported-PointMaze (blue indicates small values, red indicates large values). The state marked by the yellow star is used as the target, and we compute the MSE between its embedding and those of all other states in the maze. Since MSE is symmetric, this visualization does not fully capture directional reachability in the asymmetric teleportation dynamics. Nevertheless, with straightening, the resulting heatmaps are significantly closer to the transition-aware distances obtained using A-star.}
\label{img:heatmap_teleport}
\end{figure*}
```
```{=latex}
\begin{figure*}[h]\center
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=\textwidth]{images/plan20_0_success_grid.pdf}
        \caption{With straightening, the agent reaches the target within given step limit.}
    \end{subfigure}
    \begin{subfigure}[t]{\textwidth}

        \includegraphics[width=\textwidth]{images/plan25_0_failure_grid.pdf}
        \caption{Without straightening, the agent gets stuck at the corner.}
    \end{subfigure}
\caption{Comparison of Planning Trajectories in Teleported-PointMaze. The frames were masked by black after reaching the target.
}
\label{img:teleport_str_comp}
\end{figure*}
```

[^1]: We choose DINOv2 because it has shown the best empirical performance for latent planning on 2D navigation tasks [@terver2026drivessuccessphysicalplanning], outperforming DINOv3 [@dinov3] and V-JEPA2 [@assran2025vjepa2selfsupervisedvideo].

[^2]: We observe severe performance degradation when training without straightening and decreasing the learning rate helps. We thus use $lr=1e-6$ for no straightening.

[^3]: This is for open-loop. If using MPC, we execute the first 5 actions (or the first chunk of actions if using a frameskip of 5).