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.
Figures & tables
Figure 1: 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.
Figure 2: 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 Section E.2 .
Figure 3: During training, we minimize the prediction loss between the predicted embedding z^t and the target embedding zt 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 z^T and the goal embedding zg .
Figure 4: Action-Space Loss Landscape. We pick one test sample from PushT with a planning horizon of 25 steps. For each coordinate (ax,ay) 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.
Figure 5: 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.
Figure 6: 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. Figures 6(b) and 6(c) use ResNet with output features z∈R14×14×8 , trained with straightening regularization. After straightening, the latent distance accurately reflects the minimum number of steps required to reach the target.
Figure 7: 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 Section E.3 .
Encoder
Config
Wall
PointMaze – UMaze
PointMaze – Medium
PushT
dim
Lcurv
Open-loop
MPC
Open-loop
MPC
Open-loop
MPC
Open-loop
MPC
DINOv2 (CLS)
1×384
✗
28.67 ± 12.68
66.67 ± 10.50
25.33 ± 0.94
82.67 ± 9.98
20.00 ± 8.16
67.50 ± 3.54
19.33 ± 8.22
28.00 ± 1.63
DINOv2 (patch) + proj
1×384
✗
28.67 ± 0.94
76.00 ± 4.90
34.67 ± 1.89
79.33 ± 2.49
18.00 ± 1.63
46.00 ± 3.27
2.00 ± 1.63
11.33 ± 3.40
DINOv2 (patch) + proj
1×384
✓
32.00 ‡ ± 7.12
77.33 ‡ ± 0.94
38.67 ± 3.40
96.00 ± 0.00
22.67 † ± 5.73
78.00 † ± 2.83
2.00 ± 1.63
8.67 ± 1.89
ResNet (from scratch)
1×384
✗
58.67 ± 4.99
68.00 ± 4.32
58.00 ± 4.90
88.67 ± 0.94
73.33 ± 6.60
89.33 ± 2.49
40.00 ± 4.32
70.00 ± 1.63
ResNet (from scratch)
1×384
✓
86.67 ± 5.25
97.33 ± 2.49
63.33 ± 1.89
97.33 ± 1.89
79.33 ‡ ± 5.25
93.33 ‡ ± 2.49
58.67 † ± 5.73
83.33 † ± 2.49
Table 1: Success Rate of 50 Test Episodes (%) using the GD planner. Values are mean ± std over three eval data seeds. The best values are bold . The shaded rows are ours while the rest is DINO-WM ( Zhou et al., 2025 ) . For global features, the straightening strength λ is selected using MPC on two held-out validation seeds ( λ=0.1 unless marked: †λ=0.01 , ‡λ=0.001 ). All spatial features use λ=0.1 .
Figure 8: 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∈R14×14×8 ) with different strengths of straightening. Our model reaches 100% success rates very quickly as shown in Figures 8(a) and 8(b) .
Figure 9: 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.
Figure 10: Success rate versus wall-clock planning time for open-loop GD and CEM planning after straightening.
Method
Config
Wall
PointMaze – UMaze
PointMaze – Medium
PushT
dim
Lcurv
GD
CEM
GD
CEM
GD
CEM
GD
CEM
DINOv2 (patch)
14×14×384
✗
52.67 ± 5.73
87.33 ± 4.99
35.33 ± 4.11
88.00 ± 1.63
40.83 ± 10.07
88.00 ± 1.63
56.00 ± 4.32
71.33 ± 7.72
DINOv2 (patch) + proj
14×14×8
✗
80.00 ± 7.12
92.00 ± 0.00
44.00 ± 7.12
75.33 ± 4.99
72.00 ± 4.32
92.67 ± 4.71
70.00 ± 1.63
71.33 ± 6.18
DINOv2 (patch) + proj
14×14×8
✓
90.67 ± 0.94
100.00 ± 0.00
94.00 ± 1.63
94.00 ± 1.63
82.67 ± 3.77
86.67 ± 1.89
77.33 ± 6.18
80.00 ± 4.32
ResNet
14×14×8
✗
1.33 ± 1.89
1.33 ± 0.94
14.67 ± 4.99
20.67 ± 0.94
18.67 ± 4.11
24.00 ± 4.32
71.33 ± 7.36
56.00 ± 0.00
ResNet
14×14×8
✓
84.67 ± 2.49
90.00 ± 5.89
64.67 ± 8.38
83.33 ± 2.49
80.67 ± 0.94
89.33 ± 6.18
70.67 ± 0.94
72.67 ± 6.60
Appendix
Table 5: Open-Loop Planning Success Rate of 50 Test Trajectories (%). We compare GD and CEM planners. For CEM, we use 200 samples and 10 iterations. Values are mean ± std over three data seeds. The best value is bold .
Figure 11: Comparison of Different Dimensions. The line plots show the success rate changes with increasing channels. Too small dimensions (e.g. dv=2 ) are unable to encode sufficient planning-relevant information, while unnecessarily high dimensions (e.g. dv=128 ) hinder the planning performance.
Figure 12: Comparison with Other Regularizations.
Figure 13: Aggregation Head for Straightening during Training. The prediction loss is applied to the spatial features, while the curvature loss is applied to the aggregated features.
Figure 14: 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.
Figure 15: Distance heatmaps of PointMaze-UMaze.
Figure 16: Distance heatmaps of PointMaze-Medium.
Figure 17: PCA of Trajectories of Wall.
Figure 18: PCA of Trajectories of PointMaze-UMaze.
Figure 19: PCA of Trajectories of PointMaze-Medium.
Figure 20: PCA of Trajectories of PushT. The overlaid figures only include five samples for readability.
Figure 21: 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.
Figure 22: 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.
Figure 23: 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.
Figure 24: Open-Loop Planning Trajectories of PushT. The first row is from the simulator and the second from the decoder.
Figure 25: Teleported-PointMaze. Note that the teleportation happens inside the red box.
Figure 26: 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.
Figure 27: Comparison of Planning Trajectories in Teleported-PointMaze. The frames were masked by black after reaching the target.
Joint-embedding predictive architectures enable planning with latent world models, but accurate transition prediction alone does not ensure that the planning objective is easy to optimize. We introduce Control-Geometry Straightening (CGS), a single auxiliary loss that learns planner-friendly representations by directly straightening control geometry for sampling-efficient planning. CGS matches pairwise cosine similarities among actions to those among corresponding latent differences only using local transitions from pixel-action pairs. The loss can be applied across world-model architectures using end-to-end learned or pretrained representations. Under linear-dynamics, our theoretical analysis connects this objective to temporal straightening and more balanced terminal-cost curvature across the full planning horizon, yielding finite-budget guarantees for MPPI, local contraction results for CEM, and convergence bounds for gradient descent. Across four control environments and multiple planners, CGS improves planning with fewer sampled candidates and refinement steps, achieving success-rate gains up to 20 and 12.6 percentage points over LeWorldModel (LeWM) and its temporal-straightening variant (LeWM+TS), respectively, with sampling-based planners using 128 candidates per update. Probes, comparisons with DINO-WM architecture, and planner-side ablations clarify how latent motion organization, state dependence, and dynamical context shape planning behavior. Straightening control geometry thus makes good action sequences easier to find under limited planning budgets.
Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting in representation space rather than reconstructing pixels, making them a natural backbone for latent model predictive control from offline demonstration logs. JEPA-style training optimizes short-horizon latent prediction, whereas planning requires a multi-step ranking of imagined futures by goal progress. Prior JEPA planners often inherit that ranking from embedding geometry, typically latent Euclidean distance, which arises as a byproduct of representation learning rather than as a progress cost mined from the logs. We propose Temporal-Distance-JEPA, which retains the LeWM encoder--predictor backbone and mines a directed temporal cost from reward-free trajectories: same-trajectory step order supplies positive targets, cross-trajectory pairs act as heuristic negatives, and a rollout-consistency term matches the planner horizon. The mined supervision serves two roles: as the deployed planning cost when progress is topological, and as a representation signal that improves Euclidean planning when contact geometry dominates. Under locked evaluation, deploying the mined cost raises Two-Room success to 100.0% versus LeWM's 97.4%, while shared Euclidean planning on the same temporally trained checkpoint raises OGB-Cube by 14.2 points over LeWM and improves Push-T. Against LeWM and the concurrent RC-aux baseline under locked evaluation, Temporal-Distance-JEPA matches or exceeds both methods on every environment. Ablations show that the directed head, cross-trajectory negatives, and rollout consistency each contribute. Temporal-Distance-JEPA narrows the train--plan gap for JEPA world-model planners by discovering temporal progress structure in offline logs and co-designing cost form with plan-time deployment. Code is available at https://github.com/HKBU-KnowComp/Temporal-Distance-JEPA.
Latent world models plan by rolling a frozen predictor forward under candidate action sequences and ranking the candidates by the latent distance between their imagined end state and the goal. However, this ranking breaks down when the goal lies several plans away, because the latent distance measures how closely an end state resembles the goal rather than how far it remains from reaching it. To address this, we propose TEMPO, a temporal-distance planning objective that leaves the world model untouched, learns only from the recorded trajectories already used to train it, and adds negligible cost to the planner's search. TEMPO learns a small map of the frozen latent in which the distance between two states of an episode reflects the number of environment steps between them, and blends this distance into the planner's cost. It requires no rewards, policies or success labels and, being a cost rather than a model, applies to frozen world models with one latent vector per state that plan by a latent distance. We evaluate TEMPO on eleven simulated environments (e.g., maze navigation, tabletop pushing, robotic arm control and three-dimensional manipulation) with the LeWM and PLDM planners. With a small MLP that adds at most 0.3% to a plan's arithmetic, TEMPO improves both planners at every goal distance, including the one-plan setting of their evaluations, raises LeWM from 36% to 99% on TwoRoom three plans from the goal, and remains competitive on a broad range of 2D and 3D navigation, reaching and manipulation tasks.