PaperlayerView exact revision on arXiv

Source-linked AI summary

Temporal Straightening for Latent Planning

Ying Wang, Oumayma Bounou, Gaoyue Zhou, Randall Balestriero, Tim G. J. Rudner, Yann LeCun, Mengye Ren

arXiv:2603.12231v3cs.LG

TL;DR

Latent planning needs representations whose geometry reflects dynamics and supports reliable optimization, but pretrained visual embeddings can produce misleading distances and difficult objectives. The paper jointly trains a JEPA encoder and predictor with temporal curvature regularization, showing that straightening improves distance fidelity, planning conditioning, and goal-reaching success. The approach is evaluated primarily for continuous goal-conditioned planning with symmetric Euclidean costs.

  • Problem

    Pretrained visual encoders are not tailored to planning, so curved latent trajectories can make Euclidean goal distances misleading and planning objectives difficult to optimize.

  • Method

    The method jointly trains a JEPA encoder and predictor while regularizing local latent-trajectory curvature to produce straighter representations.

  • Results

    Across goal-reaching tasks, open-loop planning success improves by 20–60% and MPC success by 20–30% with a simple gradient-based planner.

  • Takeaways & Limitations

    Temporal straightening makes Euclidean distance a more reliable geodesic proxy and improves the conditioning and practical performance of gradient-based latent planning.

  • Takeaways & Limitations

    The formulation focuses on continuous goal-conditioned planning with a symmetric Euclidean goal cost, which may be suboptimal for asymmetric or irreversible dynamics.

Abstract

from arXiv · show

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.

1. Introduction

Latent world models make planning efficient by optimizing actions in a learned representation, but pretrained features can produce curved trajectories and poorly conditioned objectives. Temporal straightening addresses this by regularizing representations so distances better support gradient-based planning, improving goal-reaching success.

  • Latent world models compress observations, learn predictive dynamics, and optimize action sequences using latent-space goal costs.This abstraction reduces dimensionality and ignores noise compared with pixel- or state-space planning.
  • Highly curved latent trajectories make Euclidean distances misleading and can produce non-convex planning objectives that hinder gradient-based optimization.Search-based methods such as CEM and MPPI can mitigate this difficulty but add computation and latency.
  • Useful planning representations should preserve predictable dynamics while removing task-irrelevant information and regularizing distances to reflect geodesic structure.Such representations could enable differentiable planning without computationally expensive search.
  • Temporal straightening jointly trains a JEPA encoder and predictor with an explicit curvature regularizer on latent trajectories.The prediction objective induces some straightening, while explicit regularization strengthens and stabilizes it.
  • Open-loop planning success improves by 20–60% and MPC success by 20–30% with a simple gradient-based planner.

2. Related Work

Prior latent world-model methods learn representations through reconstruction, pretrained visual features, JEPA prediction, temporal contrastive learning, or straightening-related objectives. These approaches differ in how they preserve task-relevant dynamics while handling low-level detail and temporal geometry.

  • Latent world models rely on learned representations to predict dynamics and plan from compact sensory inputs.
  • Reconstruction objectives can overemphasize unnecessary low-level visual details and may miss information relevant to planning.
  • JEPA learns representations through prediction, aiming to retain predictable structure without preserving unpredictable low-level details.
  • Temporal contrastive learning encourages nearby frames to have similar embeddings and distant frames to have dissimilar embeddings, but sample selection requires careful tuning.
  • Prior work has studied implicit straightening, curvature regularization, and straightening objectives for visual and video representations.

3. Temporal Straightening

The proposed JEPA world model maps observations and actions into latent representations, predicts future latent states, and jointly trains prediction with a curvature-based straightening objective. Stop-gradient prevents collapse while the combined loss controls the strength of straightening.

  • The world model maps high-dimensional observations into a latent space and uses that space to model dynamics for planning.
  • The architecture contains a sensory encoder, an action encoder, and a predictor for future latent states.
  • Temporal straightening penalizes curvature across three consecutive latent representations by encouraging their adjacent velocity vectors to align.
  • The encoder and predictor are trained jointly to minimize prediction error while enforcing straighter trajectories.
  • Stop-gradient is used to prevent collapse of the trainable latent representation, and the curvature regularizer can be combined with other anti-collapse methods.
  • The total objective is Ltotal = Lpred + λLcurv, where λ controls the strength of straightening.

4. Planning with Straightened Dynamics

The analysis connects straighter latent dynamics to better-conditioned gradient-based planning, while experiments show reduced non-convexity in the action-space objective. The theoretical guarantee is established for linear dynamics, with nonlinear extensions left for future work.

  • The paper analyzes how straightened latent dynamics affect gradient-based planning in a linear dynamical system.
  • In a PushT horizon-25 test, the action-space loss landscape is closer to convex after straightening.
  • The planning task optimizes an action sequence over a horizon to minimize terminal mean-squared error to a target latent goal.
  • Under linear dynamics, ε-straight transitions approach a reference system whose latent state evolves linearly along trajectories modified by control inputs.
  • When ε is small, the planning Hessian remains better conditioned and its effective condition number grows slowly with the horizon.
  • Better Hessian conditioning yields faster linear convergence for gradient descent on the quadratic planning objective.
  • Under nonlinear dynamics, analogous guarantees require controlling products of state-dependent Jacobians and higher-order terms.

5. Experiments

Across four environments and multiple encoder configurations, temporal straightening improves latent geometry and gradient-based planning. It makes latent distances better reflect geodesic progress and increases planning success, while long-horizon prediction remains challenging.

  • Experimental setup: The experiments evaluate the method on Wall, PointMaze UMaze, PointMaze-Medium, and PushT against DINO-WM using spatial or global visual features.The shared setup uses a frameskip of five and a ViT dynamics predictor; encoder variants include projected DINOv2 and ResNet trained from scratch.
  • Latent geometry: JEPA training alone can implicitly straighten latent trajectories, while explicit curvature regularization further strengthens and stabilizes this effect.The pretrained DINOv2 embedding space is highly curved, whereas training generally makes trajectories straighter; explicit regularization reduces curvature further.
  • Architecture and objectives: The aggregation head reduces curvature in global features and generally improves planning, especially where patch-wise similarity cannot capture complex global state changes.For long-horizon planning, the paper notes that planning objectives may operate in a projected task- and geometry-aware space rather than directly in the prediction latent space.
  • Latent geometry: Straightening makes Euclidean latent distance closely align with geodesic distance and approximate the minimum steps to a target, even from suboptimal trajectories.The result is shown in PointMaze distance heatmaps and does not simply reflect memorization of inefficient training paths.
  • Representation sufficiency: Aggressively compressed features can retain sufficient information for planning, while preserving spatial structure generally matters more than retaining channels.DINOv2 channels can be reduced from 384 to 8 without degrading performance when patch tokens are retained; higher dimensions can even hurt some environments.
  • Planning performance: Explicit straightening produces more than a 10% boost in open-loop and MPC success rates across many setups.The method reaches 100% MPC success on Wall and UMaze within a few steps, while PushT improves more slowly because only terminal loss is used within horizon H = 5.
  • Planning performance: With straightening, gradient descent achieves a better success–latency trade-off than CEM, although CEM generally attains higher success rates at substantially longer runtimes.Long-horizon success rates still drop substantially because prediction errors compound and decoded rollouts can drift from simulator dynamics.

6. Conclusion

Temporal straightening produces latent representations that facilitate planning by making Euclidean distance a more reliable geodesic proxy and improving gradient-based planning conditioning. Across 2D goal-reaching tasks, it yields significant and consistent gains over baselines.

  • Temporal straightening makes Euclidean distance a more reliable proxy for geodesic distance in the learned representation space.
  • Straightened representations better condition gradient-based planning.
  • Across a range of 2D goal-reaching tasks, temporal straightening produces significant and consistent gains over baselines.
  • The findings identify representation geometry as important for latent planning and present trajectory straightening as a simple way to improve it.

Impact Statement

The paper aims to advance machine learning through world models with improved planning capabilities. Such systems may offer beneficial applications while also creating safety risks if deployed without adequate safeguards.

  • The work aims to advance the field of Machine Learning.
  • World models with improved planning capabilities could support robotics, autonomous systems, and scientific discovery.
  • Deploying such systems without adequate safety measures may involve potential risks.

A. Data and Environments

The experiments cover three 2D navigation settings: Wall, PointMaze with UMaze and Medium-Maze, and PushT. They vary in geometry, dynamics, interaction demands, and training data scale.

  • A.1. Wall: Wall is a two-room navigation environment separated by a wall with one narrow door.
  • A.1. Wall: Wall uses 2D displacement actions and training data comprising 1,920 trajectories of 50 time steps each.
  • A.2. PointMaze (UMaze and Medium-Maze): PointMaze uses MuJoCo dynamics and evaluates UMaze and Medium-Maze with forces applied along the x and y axes.
  • A.2. PointMaze (UMaze and Medium-Maze): The PointMaze training sets contain 2,000 UMaze trajectories and 4,000 Medium-Maze trajectories, each 100 time steps long.
  • A.3. PushT: PushT is a contact-rich task in which a pusher must move itself and a T-shaped block to a target configuration.

B. Experiments

The experiments use model-predictive control with gradient descent to optimize action sequences in latent space. States and goals are encoded, predicted trajectories are evaluated by latent MSE, and optimized actions are executed before replanning.

  • B.1. Model Predictive Control (MPC): Unlike DINO-WM, the major experiments use gradient descent rather than CEM as the MPC subplanner.
  • B.1. Model Predictive Control (MPC): The current and goal RGB observations are encoded into latent state representations using the trained encoder.
  • B.1. Model Predictive Control (MPC): An action sequence of horizon T is initialized by sampling actions from a Gaussian distribution.
  • B.1. Model Predictive Control (MPC): The planning objective minimizes MSE between the predicted final latent state and the goal state.
  • B.1. Model Predictive Control (MPC): Predicted latent trajectories are generated recursively by applying the world model to the previous latent state and action.
  • B.1. Model Predictive Control (MPC): Actions are iteratively updated using cost gradients, then the first k optimized actions are executed and the process is repeated at the next timestep.

B.3. Planning: GD vs. CEM

The paper compares gradient descent (GD) and CEM for open-loop latent planning, showing that straightening improves both while making GD more competitive in success–latency trade-offs.

  • Planner comparison: GD backpropagates through one learned rollout per optimization step, whereas CEM samples candidate action sequences and repeatedly refits a distribution to top performers.CEM uses M candidates over K iterations; GD uses N forward rollouts and N backward passes.
  • Planner comparison: 10× slower wall-clock planning makes CEM substantially more computationally expensive than GD on PushT.CEM requires at least 200 samples and 10 iterations for strong performance in the reported experiment.
  • Straightening effects: Straightening consistently improves success rates for both GD and CEM across environments and model architectures.The comparison is reported in Table 5 for 50 test trajectories.
  • Straightening effects: Straightening largely reduces the performance gap between GD and CEM, enabling GD to achieve a better success–latency trade-off.CEM often retains higher absolute success rates, but at substantially higher computational cost.

B.5. Comparison to Smoothness and Temporal Contrastive Objectives

The paper compares straightening with smoothness and temporal contrastive regularization, finding that these alternatives do not improve the reported PushT planning setup and can hurt when weighted strongly.

  • Alternative objectives: Smoothness penalizes large temporal jumps, but overly strong penalties can collapse embeddings of different states toward similar values.The collapse creates degenerate representations rather than preserving useful state distinctions.
  • Alternative objectives: Time contrastiveness treats nearby frames as positives and other frames in the same trajectory as negatives using an InfoNCE objective.Its design encourages temporal proximity in embedding space and separates temporally distant frames.
  • Alternative objectives: Suboptimal trajectories can make temporal distance disagree with geodesic distance, causing temporal contrastive learning to separate geodesically close states.This is a specific limitation of using trajectory time as the positive-negative signal.
  • Empirical comparison: No improvements are observed from smoothness or temporal contrastive objectives on open-loop GD success for PushT.The evaluation covers 50 test episodes across three data seeds; smaller weights are less harmful but still trail straightening.
  • Straightening variants: A learnable aggregation head gives the best performance among the tested straightening strategies.The variants are evaluated using projected DINOv2 patch features and open-loop planning success.

C. Theoretical Analysis

The theoretical analysis connects straighter latent dynamics to better-conditioned planning objectives, while experiments and visualizations examine distance fidelity, trajectory curvature, and representation geometry across tasks.

  • Planning Hessian: The terminal planning objective is a quadratic function of the action sequence under linear latent dynamics, with Hessian H = 2J_Φ^T J_Φ.The Hessian is positive semidefinite and its nonzero eigenvalues correspond to those of the finite-horizon controllability Gramian.
  • Scope of theory: When action dimension da < d, the controllability Gramian and planning Hessian may be singular outside the controllable subspace.The analysis then uses an effective condition number and requires additional controllability assumptions for lower bounds.
  • Straightness proxy: Cosine similarity between consecutive latent velocities serves as a practical proxy for reducing deviation of A from I along visited directions.The proxy relies on bounded velocity magnitudes and smooth actions.
  • Scope of theory: Uniform spectral control requires a coverage condition ensuring that visited velocity directions span the latent space.Without this condition, the proposition controls A − I only along directions encountered in the data.
  • Relation to prior approaches: Temporal straightening targets latent trajectory curvature rather than imposing locally or globally linear dynamics.The discussion distinguishes its geometric objective from local linearization and Koopman-style constraints on dynamics.
  • Empirical geometry: Straightening makes latent Euclidean distances align with maze geodesic distances and produces smoother PCA trajectories than DINO CLS embeddings.The visualizations compare distance heatmaps against A-star distances and plot sampled trajectories in two dimensions.
Loading 2603.12231v3…