🤖 AI

Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

Traj-LeWM enhances the lightweight visual world model LeWM by introducing a goal-conditioned latent trajectory cost that complements endpoint distance with path-aware information during both training and planning, resulting in significant performance improvements across multiple robotic manipulation benchmarks.

Original authors: Xiaodi Huang, Ziyi Ding, Jingtian Wan, Yuchen Liu, Yuan Zhang, Xiao-Ping Zhang, Jiayu Chen, Zhang Zhang, Tao Huang

Published 2026-08-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiaodi Huang, Ziyi Ding, Jingtian Wan, Yuchen Liu, Yuan Zhang, Xiao-Ping Zhang, Jiayu Chen, Zhang Zhang, Tao Huang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot how to navigate a maze or push a block to a specific spot. In the world of artificial intelligence, this is often done using something called a "world model." Think of a world model as a robot's internal daydreamer. Instead of just reacting to what it sees right now, the robot uses this daydreamer to simulate the future: "If I move my arm left, what will the world look like in one second? What about two seconds?"

For a long time, scientists have been trying to make these daydreams more accurate. A recent approach called LeWM (Lightweight World Model) was a big step forward. It learned to predict the future directly from camera pixels, skipping the need for complex 3D maps. However, even LeWM had a blind spot. It was like a GPS that only cared about your final destination. It would look at a list of possible routes and pick the one that looked like it would end up closest to the goal, without really checking if the road in the middle was full of potholes, dead ends, or traffic jams. If two routes ended at the same spot, the robot couldn't tell which one was actually safer or smoother to drive. This paper, Traj-LeWM, tries to fix that blind spot by teaching the robot to care about the whole journey, not just the finish line.

The Problem: The "Finish Line" Trap

The authors of this paper noticed a funny problem with how robots like LeWM make decisions. Imagine you are playing a video game where you have to guide a character to a treasure chest. You have two paths to choose from:

Path A: A straight, clear road that leads directly to the chest.
Path B: A winding, dangerous path that goes through a trap-filled forest but also ends up right next to the chest.

Old-school AI (like LeWM) looks at the end of both paths, sees they are equally close to the chest, and picks one at random. It doesn't realize that Path B might crash the character into a tree halfway through. The paper argues that this happens because the robot only checks the "endpoint distance." It ignores the "intermediate path"—the messy, tricky stuff that happens in the middle of the action.

The researchers found that in complex tasks, like pushing a T-shaped object or navigating a room with obstacles, two different action sequences could end up in the exact same spot, but one would succeed while the other would fail spectacularly. The old robot couldn't tell the difference because it was too focused on the finish line.

The Solution: The "Whole Journey" Scorecard

To fix this, the team introduced Traj-LeWM. They kept the robot's ability to predict the future but added a new "scorecard" called the Latent Trajectory Cost (LTC).

Think of LTC as a strict coach who doesn't just look at where the athlete finishes, but watches the entire race.

The Old Way: The coach says, "Great job! You finished 1 meter from the finish line."
The New Way (Traj-LeWM): The coach says, "You finished 1 meter from the line, but you stumbled three times, ran into a wall, and took a weird detour. That's a bad run. Let's try the other runner who finished at the same spot but ran smoothly."

The new system learns a "cost" for every possible path. If a path looks like it will hit a wall or take a weird turn, the cost goes up. If the path is smooth and logical, the cost stays low. The robot then combines this "smoothness score" with the old "distance to goal" score to pick the best action.

How They Taught the Robot

Teaching a robot to care about the middle of the journey is tricky. You can't just tell it "don't hit walls" because it doesn't know what a wall is yet. So, the researchers used a clever trick called preference learning.

They showed the robot pairs of stories:

The Good Story: A perfect path where the robot successfully pushes the block to the goal.
The Bad Story: A path that looks similar but has a mistake (like hitting a wall) or a path that goes to the wrong goal entirely.

The robot was asked to guess which story was better. By getting thousands of these "this one is better than that one" lessons, the robot learned to assign a low cost to good paths and a high cost to bad ones. They even let the robot try to plan in a simulated environment, and if it crashed, they saved that crash as a "bad story" to learn from. This is like a video game player learning not to jump off a cliff because they died doing it, but this time the robot is learning from its own simulated mistakes.

What They Found

The team tested this new robot in four different simulated worlds:

Push-T: Pushing a T-shaped object to a target.
OGBench-Cube: Moving a cube to a specific spot.
Reacher: A robotic arm reaching for a target.
Two-Room: Navigating through a maze with two rooms.

In these simulations, the new Traj-LeWM robot was significantly better than the old LeWM robot.

On the Cube task, it improved success rates by 14 percentage points (going from 74% to 88%).
On Reacher, it improved by 7 percentage points (from 86% to 93%).
On Two-Room, it improved by 7 percentage points (from 87% to 94%).
On Push-T, it improved by 3 percentage points (from 96% to 99%).

They also tested it on a real physical robot (a Franka FR3 arm) in the real world. In 20 real-life tasks, the old robot succeeded 10 times (50%), while the new Traj-LeWM robot succeeded 14 times (70%). While this is a small test, it suggests the method works outside of the computer simulation, too.

Why This Matters

The paper shows that for robots to be truly smart, they can't just look at the destination; they have to respect the journey. By teaching the robot to evaluate the entire path it plans to take, rather than just the final spot, it becomes much better at avoiding traps and navigating complex environments.

The researchers confirmed that this improvement comes from two things:

Better Learning: The robot learned a better internal map of the world because it was trained to care about the whole path, not just the end.
Better Planning: Even when the robot's prediction of the future was slightly wrong, the "whole journey" score helped it pick the safer option.

In short, Traj-LeWM proves that in the world of robot planning, it's not just about where you end up; it's about how you get there. By giving the robot a way to appreciate a smooth, safe path over a bumpy, risky one, the authors have taken a significant step toward making robots that can handle the messy, unpredictable real world.

Technical Summary of Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

1. Problem Statement

The paper addresses two specific limitations in LeWM (Lightweight Visual World Model), a joint-embedding predictive architecture (JEPA) that learns latent dynamics from pixels to perform goal-conditioned planning. While LeWM achieves strong performance on simple tasks, its success rate drops significantly on more complex environments (e.g., from 96% on Push-T to 74% on OGBench-Cube). The authors identify two root causes for this degradation:

Missing Trajectory-Level Supervision: LeWM's training objective relies on next-step prediction (local transitions) and a regularization term (SIGReg) to prevent representation collapse. Neither component explicitly requires the shared encoder to preserve information about the complete trajectory relative to the task goal. Consequently, the latent space is not directly trained to assess the quality of a full candidate trajectory.
Endpoint-Only Scoring Insensitivity: During planning, LeWM ranks candidate action sequences solely based on the Euclidean distance between the predicted latent endpoint and the goal representation. This approach ignores intermediate path information. Two trajectories with identical endpoints but different intermediate evolutions (one successful, one failing) receive the same score, leading to suboptimal candidate selection when model predictions diverge from actual execution outcomes.

2. Methodology: Traj-LeWM

The authors propose Traj-LeWM, a framework that retains LeWM's original local-dynamics objective and endpoint scoring while introducing a Goal-Conditioned Latent Trajectory Cost (LTC). The method operates in three stages:

A. Goal-Conditioned Latent Trajectory Cost (LTC)
The core innovation is a learnable scalar cost function, $S_\psi(\tau | g)$ , which evaluates a complete latent trajectory $\tau$ given a goal $g$ . Unlike endpoint distance, LTC aggregates information over the entire path. It is defined as:
$S_\psi(\tau | g) = \frac{1}{T} \sum_{t=0}^{T-1} \ell_\psi(z_t, \Delta z_t, d^g_t, \eta_t)$
Where:

$z_t$ : Current latent state.
$\Delta z_t$ : Local evolution (difference between adjacent states).
$d^g_t = z_g - z_t$ : Goal-relative quantity (conditioning).
$\eta_t$ : Normalized temporal phase.
$\ell_\psi$ : A learnable non-negative per-step contribution function.

This structure allows the cost to vary based on how a state is reached, not just where it ends.

B. Trajectory-Preference Supervision
To train the LTC and shape the shared encoder, the authors introduce a preference-based learning mechanism using pairwise logistic loss. The system learns to assign lower costs to "preferred" (expert-matched) trajectories and higher costs to "negative" trajectories. Two types of negative examples are used:

Synthetic Preferences:
- Goal Mismatch: An expert trajectory paired with a shifted, unrelated goal.
- Endpoint-Preserving Perturbation: An expert trajectory with noise added to intermediate states (while keeping start/end fixed). This forces the model to distinguish between different intermediate paths even if endpoints are identical.
Closed-Loop Failure Preferences:
- The current model performs planning using only endpoint scoring (without LTC) and executes the plan in the environment.
- Failed executions are mined and paired with expert trajectories having the same start and goal.
- These real-world failures provide negative examples that the endpoint-only metric failed to identify, calibrating the LTC against the model's specific failure modes.

C. Joint Training and Planning

Training: The total loss combines the original LeWM prediction loss ( $L_{LeWM}$ ), the synthetic preference loss ( $L_{path}$ ), and the mined failure loss ( $L_{mined}$ ). The synthetic branch updates both the LTC and the shared encoder, while the mined branch updates only the LTC (via detached gradients) to avoid confounding the predictor's dynamics learning.
Planning: Candidates are ranked using a Joint Score:
$C_\lambda(\hat{\tau}, g) = C_0(\hat{\tau}, g) + \lambda \tilde{S}_\psi(\hat{\tau} | g)$
Where $C_0$ is the original endpoint distance, and $\tilde{S}_\psi$ is a calibrated LTC score. This allows the planner to utilize intermediate-path information when endpoint predictions are ambiguous or inaccurate.

3. Key Contributions

Latent Trajectory Cost (LTC): Introduction of a learned, goal-conditioned cost that evaluates complete predicted trajectories, aggregating goal-relative latent evolution to provide a path-sensitive planning signal.
Trajectory-Preference Learning: Development of a mechanism using synthetic negatives (goal mismatches and intermediate perturbations) and mined closed-loop failures to train the LTC and shape the shared representation without altering the original predictor objective.
Comprehensive Evaluation: Extensive testing on four simulated tasks (Push-T, OGBench-Cube, Reacher, Two-Room) and a physical Franka FR3 robot, demonstrating improvements over the baseline.

4. Experimental Results

Simulation Results:
Traj-LeWM outperforms LeWM across all four tasks, with success rate improvements of:

Push-T: +3% (96% $\to$ 99%)
OGBench-Cube: +14% (74% $\to$ 88%)
Reacher: +7% (86% $\to$ 93%)
Two-Room: +7% (87% $\to$ 94%)

Real-Robot Evaluation:
On a physical Franka FR3 with 20 start-goal tasks, Traj-LeWM achieved a 70% success rate (14/20) compared to LeWM's 50% (10/20). Failures due to collisions decreased from 20% to 10%.

Controlled Experiments & Ablations:

Representation Shaping (Gap 1): Even when LTC is removed at test time, Traj-LeWM's shared encoder (trained with trajectory preferences) produces endpoint scores with higher AUC for distinguishing success vs. failure compared to LeWM.
Path Awareness (Gap 2): When endpoint predictions are inaccurate, adding LTC improves the correct ordering of candidates. In "endpoint-controlled" tests (where start and end states are identical), LTC successfully distinguishes successful from failed trajectories in >94% of cases on Reacher, proving it utilizes intermediate path information.
Complementarity: Ablation studies show that trajectory supervision (training) and explicit LTC ranking (planning) are complementary. The joint score yields the best results, indicating they are not interchangeable.

5. Significance and Claims

The paper claims that Traj-LeWM successfully addresses the limitations of local-dynamics-only world models by integrating trajectory-level information into both representation learning and candidate ranking.

Modest Claims: The authors explicitly state that the real-robot results are "preliminary evidence of real-robot feasibility rather than a conclusive performance advantage" due to the limited number of paired tasks (20).
Core Insight: The work demonstrates that endpoint distance alone is insufficient for reliable planning in complex environments. By learning a cost over complete latent rollouts, the model can leverage intermediate-path evidence to rank candidates more effectively, particularly when endpoint predictions are uncertain.
Efficiency: The method adds fewer than 1M parameters to the model, maintaining a lightweight architecture comparable to the original LeWM while significantly boosting performance on challenging manipulation and navigation tasks.