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. 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.

• The paper introduces temporal straightening to align Euclidean distances with true geodesic costs in latent spaces, enhancing planning reliability.
• It integrates a curvature regularizer into the predictive loss, boosting gradient-based optimization with up to 94–100% MPC success in complex tasks.
• Empirical analysis on 2D navigation and contact-rich environments shows that straightened latent representations reduce reliance on expensive optimization methods.

Temporal Straightening for Latent Planning: An Expert Essay

Motivation and Problem Analysis

Planning in high-dimensional environments with learned world models remains a formidable challenge due to the geometry of the latent spaces produced by visual encoders. Although recent world model architectures based on predictive modeling (e.g., JEPA, DINO-WM) have demonstrated efficiency and generalization by abstracting high-dimensional sensory streams into compact representations, the geometry of the induced latent spaces is not naturally suited for direct optimization. Latent trajectories in these spaces are often highly curved, leading to misalignments between Euclidean and true geodesic distances, complicating planning objectives and the convergence of gradient-based optimizers. The utility of pretrained semantic-rich visual representations (such as DINOv2) is limited by this geometric misalignment, resulting in suboptimal planning, particularly when gradient-based methods are used.

Temporal Straightening: Methodology

This paper introduces temporal straightening as a geometric regularization for latent planning. Drawing on the perceptual straightening phenomenon in human vision, the authors propose augmenting the representation learning pipeline with a curvature regularizer. This regularizer directly penalizes latent trajectory curvature by maximizing the local cosine similarity of successive latent velocity vectors: $$\mathcal{C} = \frac{v_t \cdot v_{t+1}}{||v_t||_2 \cdot ||v_{t+1}||_2}$$ where $v_t = z_{t+1} - z_t$ and latent representations $z_t$ are learned through a differentiable encoder.

The regularization is combined with a standard predictive loss in a joint objective: $$\mathcal{L}_{\text{total}} = \mathcal{L}_{pred} + \lambda \mathcal{L}_{curv}$$ with stop-gradient employed on the target branch for collapse prevention. This approach constrains the latent dynamics such that Euclidean distances become more congruent with the geodesic structure of feasible transitions.

Figure 1: Latent trajectories before and after straightening, illustrating the reduction of curvature and more faithful metric alignment for planning.

Theoretical Justification

The effect of temporal straightening is formalized via analysis of the planning loss landscape under linearized latent dynamics: $z_{t+1} = A z_t + B a_t$ The authors prove that straightening (forcing $A \approx I$) substantially improves the condition number of the planning Hessian. Specifically, for $\varepsilon$-straight dynamics ($\|A - I\|_2 \leq \varepsilon$), the effective condition number of the planning objective grows as $e^{O(\varepsilon K)}$ with planning horizon $K$, ensuring more stable and rapid convergence for gradient-based optimizers. Cosine similarity of successive latent velocities is shown to act as a practical regularization proxy, under realistic assumptions on action smoothness and latent norm stability.

Architectural Implementation

The architecture comprises (1) a visual encoder (either a trainable projector atop a frozen DINOv2 backbone or a trainable ResNet from scratch), (2) an action encoder, and (3) a Transformer-based predictor. For spatial latents, straightening operates either patch-wise or via a learned pooling head, with empirical ablations showing superior performance for the latter, reflecting the benefit of promoting straightness at the level of global dynamics rather than per-patch changes.

Figure 2: Training and planning framework integrating curvature regularization in the objective.

Empirical Analysis

Effect on Latent Space Geometry

Quantitative results confirm that pretrained DINOv2 features without regularization induce highly curved latent trajectories (low cosine similarity). Both implicit and explicit straightening dramatically increase cosine similarity, resulting in smoother trajectories and better metric alignment.

Figure 3: Reduced curvature directly correlates with improved open-loop planning success rates.

Visualization of latent distance heatmaps show that post-straightening, the Euclidean metric in embedding space more accurately reflects minimum-step geodesic costs, even without expert trajectory supervision. Importantly, both local and long-range structures become more consistent, supporting global planning objectives.

Planning Performance

Substantial empirical gains are observed in goal-reaching success rates (open-loop and closed-loop). On challenging 2D navigation and contact-rich tasks (Wall, PointMaze-UMaze/Medium, PushT), the following effects are robustly repeated:

• Gradient Descent optimization achieves 20–60% higher open-loop and 20–30% higher MPC success rates with straightening regularization.
• Explicit straightening yields up to 94–100% MPC success on certain mazes within a handful of steps, outperforming DINO-WM and variants using CEM.
• Preserving spatial structure in the encoder is far more important than maintaining high channel dimensionality; aggressive channel compression is possible without loss of planning efficacy.

Long-horizon planning (e.g., 50 steps in PushT) exposes typical compounding model errors but straightened latents consistently outperform baselines, evidencing improved robustness in more demanding planning regimes.

Figure 4: Long-horizon planning in PushT visualizes robust performance in the presence of straightened latent geometry.

Robustness and Generalization

To disambiguate appearance versus dynamics, a teleportation variant of PointMaze is introduced. In this setting, DINOv2’s visual similarity-derived embeddings catastrophically fail, whereas straightened representations enable successful exploitation of dynamical structure (i.e., teleport). This highlights that straightened latents do not simply encode visual similarity, but track controllable environment dynamics.

Implications and Future Directions

The main implication is that the geometry of latent spaces is a primary lever for controllable planning difficulty in world model-based reinforcement learning. Temporal straightening regularization yields planning-suitable representations that both improve gradient-based solver efficiency and reduce reliance on expensive population-based optimizers (such as CEM or MPPI), with strong numerical evidence.

Practically, this method is simple to implement atop existing predictive architectures, requiring only modification of the objective by a differentiable regularizer and, optionally, a pooling head. Theoretically, the work connects the geometry of learned representations, metric faithfulness, and optimization conditioning, suggesting new approaches for constructing world models tailored to algorithmic downstream consumers.

A natural future direction is extending the approach to more complex, high-dimensional, partially observable, or stochastic environments, and examining the interaction between model class, regularizer strength, and sample complexity. Open problems remain for guaranteeing metric alignment under non-linear or high-variance dynamics and for extending the straightening concept to episodic or hierarchical settings.

Conclusion

Temporal straightening is demonstrated to be a robust and general technique for equipping latent world models with planning-suitable geometric properties. By aligning Euclidean distances in the latent space with environment geodesics and improving objective conditioning, it materially enhances the performance and reliability of gradient-based planners across a suite of goal-reaching tasks. This technique represents a necessary refinement in the interplay between representation learning and control in visual model-based RL, and provides actionable guidance for the development of future latent world models targeting efficient policy optimization (2603.12231).