home research people contact donate

Temporal Straightening for Latent Planning

Ying Wang1, Oumayma Bounou1, Gaoyue Zhou1, Randall Balestriero2, Tim G. J. Rudner3, Yann LeCun1*, and Mengye Ren1*

1New York University, 2Brown University, 3University of Toronto
*Equal advising

The 43rd International Conference on Machine Learning (ICML 2026)

arXiv PDF Code

TL;DR: Inspired by the perceptual straightening hypothesis in human vision, we introduce temporal straightening to improve representation learning 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. 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.

Overview

Method

Inspired by the perceptual straightening hypothesis in human vision, which posits that visual systems transform complex 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 world model, while imposing regularization on the curvature of latent trajectories during training. The training objective is:

$$\mathcal{L}_{\text{pred}} = \lVert \hat{z}_{t+1} - \mathrm{sg}(z_{t+1}) \rVert_{2}^{2}$$ $$\mathcal{L}_{\text{curv}} = 1 - C, \quad \text{where } C = \cos(z_{t+1} - z_{t},\ z_{t+2} - z_{t+1})$$ $$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{pred}} + \lambda\, \mathcal{L}_{\text{curv}}$$

Here, $\mathrm{sg}$ denotes stop-gradient and $\lambda$ controls the strength of the straightening.

Training and planning architecture.
Training and planning architecture.

How Good Is the Embedding Space?

We inspect the learned embedding space by measuring latent trajectory curvatures, PCA projections of latent trajectories, and latent Euclidean distances to understand the impact of straightening.

  1. Implicit straightening can happen when training the encoder using the predictor loss alone.
  2. Adding straightening regularization further decreases curvature of the resulting embeddings.
  3. Straightening encourages the latent Euclidean distance to better align with the geodesic distance.
  4. Near-perfect reconstruction can be attained with a very low feature dimensionality.
Latent curvature and open-loop GD success rate for different encoders. Higher cosine similarity indicates lower curvature.
Latent curvature and open-loop GD success rate for different encoders. Higher cosine similarity indicates lower curvature.

We visualize the Euclidean distance between the embedding of a target state (denoted by the star) and all other states in the maze. Blue indicates smaller distance, red indicates larger distance.

UMaze ground-truth heatmap
UMaze: ground-truth geodesic distance.
UMaze ResNet global heatmap
UMaze: ResNet-global after straightening.
UMaze DINO CLS heatmap
UMaze: DINO CLS.
UMaze DINO patch heatmap
UMaze: DINO patch.
Medium ground-truth heatmap
Medium: ground-truth geodesic distance.
Medium ResNet global heatmap
Medium: ResNet-global after straightening.
Medium DINO CLS heatmap
Medium: DINO CLS.
Medium DINO patch heatmap
Medium: DINO patch.

We also visualize the learned trajectory representations using PCA. While latent trajectories are highly curved in the pretrained embedding space, they become significantly smoother after straightening, and Euclidean distance becomes a more faithful proxy for geodesic progress toward the goal.

Wall PCA overview
Wall PCA example 1
Wall PCA example 2
Wall PCA example 3
UMaze PCA overview
UMaze PCA example 1
UMaze PCA example 2
UMaze PCA example 3
Medium PCA example 1
Medium PCA example 2
Medium PCA example 3
PushT PCA overview
PushT PCA example 1
PushT PCA example 2
PushT PCA example 3

Planning

We perform gradient-based planning using our models on four environments: Wall, PointMaze-UMaze, PointMaze-Medium, and PushT. We report both open-loop planning and closed-loop MPC. Open-loop planning optimizes a length-$H$ action sequence using the terminal embedding distance to the target, while MPC executes the first action and replans at every step. Across environments, temporal straightening substantially improves planning performance.

Main planning results table.
Main planning results table.

Closed-loop MPC replans at every step. The success-rate curves below show that straightening reaches high MPC success quickly, especially on Wall and UMaze.

Wall MPC results
Wall.
UMaze MPC results
PointMaze-UMaze.
Medium MPC results
PointMaze-Medium.
PushT MPC results
PushT.

Below are examples of MPC planning across the four environments.

BibTeX

@inproceedings{wang2026temporal,
  title     = {Temporal Straightening for Latent Planning},
  author    = {Wang, Ying and Bounou, Oumayma and Zhou, Gaoyue and Balestriero, Randall and Rudner, Tim G. J. and LeCun, Yann and Ren, Mengye},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2026}
}