PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.14125v1cs.AI

TL;DR

LeWM lacks explicit goal-conditioned trajectory supervision and relies on endpoint distance, although candidates with similar endpoint costs can execute differently. Traj-LeWM adds a latent trajectory cost to representation learning and candidate ranking, improving success over LeWM by 3–14 percentage points across four simulated tasks.

  • Problem

    LeWM lacks explicit goal-conditioned trajectory supervision and endpoint distance omits intermediate-path information relevant to candidate quality.

  • Method

    Traj-LeWM learns a goal-conditioned latent trajectory cost from trajectory preferences and combines it with endpoint distance for representation learning and candidate ranking.

  • Results

    3, 14, 7, and 7 percentage-point gains over LeWM occur on Push-T, OGBench-Cube, Reacher, and Two-Room, respectively.

  • Takeaways & Limitations

    Controlled tests and ablations support complementary roles for trajectory-level representation shaping and path-aware candidate ranking.

  • Takeaways & Limitations

    The limited number of paired real-robot tasks makes the real-robot results preliminary evidence of feasibility rather than a conclusive performance advantage.

Abstract

from arXiv · show

LeWM is a lightweight visual world model that learns latent dynamics end-to-end from pixels and ranks candidate action sequences by the distance between their predicted endpoints and the goal. However, LeWM has two limitations. First, during training, it learns local next-step transitions without evaluating complete trajectories relative to the task goal. Second, during planning, it ranks candidates solely by predicted endpoint distance. Because model predictions may differ from actual execution outcomes, the candidate whose predicted endpoint is closest to the goal may not perform best when executed in the environment. The evolution of the complete predicted trajectory can therefore provide complementary information beyond endpoint distance. To address these limitations, we propose Traj-LeWM, which retains LeWM's local-dynamics objective and endpoint score while introducing a goal-conditioned latent trajectory cost (LTC) that aggregates trajectory-level information as a complementary signal. During training, LTC-based trajectory-preference supervision complements next-step prediction in shaping the shared representation. During planning, LTC is combined with endpoint distance to incorporate intermediate-path information into candidate ranking. With joint endpoint-plus-LTC scoring, Traj-LeWM outperforms LeWM on Push-T, OGBench-Cube, Reacher, and Two-Room by $3$, $14$, $7$, and $7$ percentage points, respectively. Controlled experiments and ablations further verify the complementary roles of trajectory-level representation shaping and path-aware candidate ranking.

1 Introduction

Traj-LeWM addresses LeWM’s endpoint-only limitations by learning a goal-conditioned latent trajectory cost over complete predicted trajectories and combining it with endpoint distance. This path-aware design improves simulated-task success and real-robot performance while preserving LeWM’s original prediction objective.

  • Motivation: LeWM’s success drops from 96% on Push-T to 74% on OGBench-Cube, while similar endpoint costs can yield different execution outcomes.These observations motivate incorporating complete-trajectory information into representation learning and candidate ranking.
  • Motivation: LeWM supervises local next-step transitions but does not directly train its shared representation to preserve complete trajectories relative to task goals.SIGReg regularizes latent-representation distributions, while autoregressive rollout generates predicted states without direct complete-trajectory supervision.
  • Method: Traj-LeWM introduces a goal-conditioned latent trajectory cost that aggregates goal-relative latent evolution over complete trajectories and complements endpoint distance with path-sensitive information.The method retains LeWM’s original prediction objective and endpoint score while mapping each complete latent trajectory to a learned scalar cost.
  • Method: Trajectory-preference learning uses synthetic negative trajectories and closed-loop execution failures to shape trajectory-level representations and train LTC.Synthetic preferences favor goal-matched expert trajectories, while mined failures further train LTC without changing LeWM’s original predictor objective.
  • Evaluation: 3, 14, 7, and 7 percentage points are the success-rate improvements over LeWM on Push-T, OGBench-Cube, Reacher, and Two-Room, respectively.Real-robot success also increases from 50% to 70% over the same 20 tasks; controlled analyses and ablations verify representation shaping and intermediate-path ranking effects.

2 Related Work

Related work spans predictive world models and JEPA-based latent planning, trajectory-level signals beyond endpoint distance, and preference-based cost learning. Traj-LeWM builds on these directions by combining latent-path scoring with endpoint-based candidate ranking.

  • World models and JEPA-based planning: World models learn predictive dynamics and plan through imagined rollouts, including generative latent agents, value-guided planners, and reward-free latent planners.Representative approaches include Dreamer, TD-MPC, PLDM, and DINO-WM.
  • World models and JEPA-based planning: JEPA-based methods predict masked or future content in representation space, with recent systems extending action-conditioned physical planning to image-goal robot control.V-JEPA 2-AC plans from image goals on robots, while a systematic study examines representation, prediction, and planner design.
  • Trajectory-level planning signals: Trajectory-level planning signals complement endpoint distance through goal-conditioned values, reachability estimates, dense goal-conditioned rewards, and learned reward predictions.These signals support planning and latent imagination in approaches including DreamerV3.
  • Trajectory-level planning signals: LTC scores the same CEM candidates from their complete predicted latent paths and supplements the endpoint score during ranking.It does not replace LeWM’s dynamics model or candidate generator.
  • Preference-based cost learning: Preference-based cost learning uses pairwise comparisons to learn scalar rankings, including Bradley–Terry models, human-comparison reward learning, and ranked-demonstration methods such as T-REX.The Bradley–Terry model provides a standard logistic form for pairwise preference modeling.

3 Method

Traj-LeWM extends LeWM with a goal-conditioned latent trajectory cost that evaluates complete predicted paths, shaping representations during training and supplementing endpoint distance during planning. Its joint endpoint-plus-LTC scoring addresses missing trajectory-level supervision and endpoint-only path insensitivity.

  • Overview: Traj-LeWM adds a goal-conditioned latent trajectory cost (LTC) that evaluates complete latent trajectories and supplements endpoint distance with path-sensitive candidate ranking.Trajectory-level preferences train LTC and shape the shared representation, while LTC provides complementary information during planning.
  • Motivation: LeWM’s local prediction and representation-regularization objectives provide no explicit trajectory-level, goal-conditioned supervision for the shared representation.The prediction loss constrains local transitions, while SIGReg shapes the latent distribution without jointly using an ordered complete trajectory and its task goal.
  • Motivation: Endpoint-only scoring cannot directly use intermediate dynamics, leaving trajectory-level, goal-conditioned information without an explicit role in candidate evaluation.Two candidates can share an initial state and endpoint while differing in intermediate latent-state sequences.
  • LTC design: LTC aggregates latent states, their temporal changes, temporal position, and goal-relative information into a nonnegative learned cost over the complete trajectory.Goal conditioning is introduced at every step, and contributions are averaged over the trajectory.
  • Trajectory-preference supervision: Trajectory-preference supervision trains LTC to assign lower costs to goal-matched expert trajectories than to mismatched, perturbed, or failed alternatives.Goal mismatch and endpoint-preserving perturbations provide complementary synthetic negatives, while closed-loop failure mining supplies planning-error preferences.
  • Planning: The combined score retains endpoint-based goal matching while incorporating information from complete predicted trajectories into candidate selection.C0 and Sψ are normalized using interquartile ranges estimated over the same endpoint-only CEM candidate set.

4 Experiments

Experiments evaluate Traj-LeWM across four simulated tasks, controlled tests, real-robot trials, and ablations. Joint endpoint-plus-LTC scoring improves closed-loop planning over LeWM, while trajectory supervision and planning-time LTC provide complementary benefits.

  • Experimental setup: Traj-LeWM and LeWM share the offline dataset, encoder–predictor backbone, predictor objective, and optimization schedule, isolating trajectory-preference supervision as the added training component.Closed-loop failures are collected only to construct preference pairs and are excluded from predictor transition-training data.
  • Main results: 99%, 88%, 93%, and 94% mean success rates are achieved on Push-T, Cube, Reacher, and Two-Room, respectively, yielding LeWM gains of 3, 14, 7, and 7 percentage points.Results use joint endpoint-plus-LTC scoring and average three evaluation seeds.
  • Controlled experiments: 0.4, 1.0, and 4.3 percentage-point endpoint-score AUC gains occur on Push-T, Cube, and Two-Room, while Reacher changes by nonsignificant −0.4 points.The endpoint-only comparison removes LTC from planning, isolating representation shaping from path-aware candidate ranking.
  • Controlled experiments: 94.4% of 207 endpoint-matched Reacher success–failure pairs are correctly ranked by LTC, and intermediate-state perturbations increase LTC cost in all 200 fixed-endpoint pairs.These tests assess path use while controlling endpoint information.
  • Real-robot evaluation: 14 of 20 real-robot trials succeed with Traj-LeWM versus 10 with LeWM, while target misses decrease from six to four and collision terminations from four to two.Because the paired-task sample is limited, these results are preliminary evidence of feasibility rather than a conclusive performance advantage.
  • Ablations: 89.0% average success under endpoint-only scoring rises to 93.5% with endpoint-plus-LTC scoring, while LTC-only ranking averages 80.3% and reaches 94% on Reacher.Under endpoint-only scoring, synthetic preferences raise average success from 85.8% to 87.3%, and adding closed-loop failures raises it to 89.0%.

5 Conclusion · Appendix

Traj-LeWM addresses the limitations of local prediction objectives and endpoint-only planning by learning and using a goal-conditioned Latent Trajectory Cost (LTC). It incorporates trajectory information into both shared-encoder representation learning and candidate ranking through joint endpoint-plus-LTC scoring.

  • 5 Conclusion: Traj-LeWM addresses missing trajectory-level constraints in local prediction objectives and the omission of intermediate paths in endpoint-distance planning.These are identified as two limitations of goal-conditioned world models.
  • 5 Conclusion: Traj-LeWM learns a goal-conditioned Latent Trajectory Cost from trajectory-level preferences.The LTC provides trajectory-level information for the proposed method.
  • 5 Conclusion: Trajectory information shapes representation learning through the shared encoder and supports candidate ranking during planning.The method uses LTC for both training-time representation learning and planning-time candidate ranking.
  • 5 Conclusion: Joint endpoint-plus-LTC scoring combines endpoint distance with trajectory information in planning.This scoring design addresses the limitation of relying only on predicted endpoint distance.

A Reproducibility Details · A.1 Model, Training, and Mining · A.2 Planning

The reproducibility setup specifies Traj-LeWM’s architecture, optimization and failure-mining procedures, alongside CEM-based MPC planning and task-specific joint-scoring weights. Implementation details include model dimensions, training hyperparameters, trajectory-pair construction, candidate-selection settings, and LTC calibration.

  • A.1 Model, Training, and Mining: The visual encoder uses 14 × 14 patches, hidden dimension 192, 12 Transformer blocks, 3 attention heads, and MLP dimension 768.The autoregressive predictor uses 6 Transformer blocks, 16 attention heads, head dimension 64, MLP dimension 2048, and dropout 0.1.
  • A.1 Model, Training, and Mining: The LTC applies three width-512 hidden layers with LayerNorm, GELU, dropout 0.1, a positive scalar output, and temporal averaging.LeWM has 18,042,672 parameters, while Traj-LeWM-syn and Traj-LeWM each have 18,867,505; LTC adds 824,833 parameters.
  • A.1 Model, Training, and Mining: AdamW training uses learning rate 5 × 10^-5, weight decay 10^-3, batch size 128, bf16 precision, gradient clipping 1.0, and linear-warmup cosine annealing.Trajectory supervision uses λpath = λmined = 0.05, β = 0.2, SIGReg weight 0.09, 17 knots, and 1,024 projections.
  • A.1 Model, Training, and Mining: After each epoch, failure mining evaluates 200 start–goal queries with endpoint-only CEM and stores detached latent pairs in a FIFO buffer of capacity 2,048.Rollouts failing to reach termination within 50 steps provide negative trajectories, paired with positive trajectories from dataset segments sharing the same start and goal.
  • A.2 Planning: Formal evaluation uses MPC with CEM, sampling 300 action sequences, retaining 30 elites, updating Gaussian statistics, and enabling warm starts across replanning rounds.The initial distribution has zero mean and unit variance in standardized action coordinates; later statistics are empirical elite statistics.
  • A.2 Planning: Joint planning sets wLTC = λ IQR(Cendpoint) max{IQR(CLTC), 10^-12}, with IQRs estimated from endpoint-only calibration candidates.This normalization defines the retained implementation’s LTC weighting relative to endpoint-score variability.
  • A.2 Planning: The final joint-scoring configuration uses λ = (0.5, 0.8, 0.6, 0.5) for Push-T, Cube, Reacher, and Two-Room, respectively.These task-specific values are the configuration adopted for the paper.

B Additional Results and Diagnostics · B.1 Verified Multi-Seed Evaluation · B.2 Controlled Path-Ranking Diagnostics

The additional diagnostics verify Traj-LeWM across evaluation seeds and define controlled analyses for comparing endpoint distance with LTC-based trajectory preferences. These analyses include multi-seed success-rate reporting, candidate-selector agreement, endpoint ROC–AUC, and endpoint-matched LTC ordering accuracy.

  • B.1 Verified Multi-Seed Evaluation: Table 5 reports Traj-LeWM results across evaluation seeds 42–44, using 50 start–goal queries per task and seed.Success-count entries denote successes out of 50 queries for each seed.
  • B.1 Verified Multi-Seed Evaluation: Success rates are reported as the mean ± sample standard deviation across the three evaluation seeds, in percent.The evaluation seeds are 42–44.
  • B.1 Verified Multi-Seed Evaluation: Figure 6 measures agreement between candidate selectors on shared candidate pools using exact top-1 agreement and overlap between the two top-5 sets.Bars show query means, error bars show query-level bootstrap 95% confidence intervals, and the gray band denotes the expectation for independent rankings.
  • B.2 Controlled Path-Ranking Diagnostics: Success rate is defined as the percentage of queries that terminate successfully, with each query assigned a binary success indicator.For N queries, the rate is computed from the indicators sq = 1 for success and sq = 0 otherwise.
  • B.2 Controlled Path-Ranking Diagnostics: Endpoint ROC–AUC evaluates how often successful candidates rank above failed candidates when candidates are scored by −Cendpoint.Ties contribute one half, while queries containing only successful or only failed candidates are excluded.
  • B.2 Controlled Path-Ranking Diagnostics: LTC pair-ordering accuracy scores endpoint-matched successful–failed pairs according to whether LTC assigns lower cost to the successful trajectory.Correct, incorrect, and tied orderings receive scores of one, zero, and one half, respectively; the reported percentage is 100 times the mean score.
  • B.2 Controlled Path-Ranking Diagnostics: Endpoint-matched executed trajectories pair each successful trajectory with an unused failed trajectory having the nearest endpoint energy, subject to a normalized endpoint-gap caliper of 0.25.Table 6 reports the percentage of pairs where LTC assigns lower cost to the successful trajectory.
  • B.2 Controlled Path-Ranking Diagnostics: Table 6 reports LTC ordering accuracy in percent for endpoint-matched successful and failed executed trajectories.Confidence intervals are omitted because the corrected aggregate count is insufficient to reconstruct the previous query-level bootstrap interval.

B.3 Candidate-Selector Agreement · B.4 Open-Loop Latent Rollout Consistency · C Real-Robot Evaluation

LTC produces candidate rankings substantially different from endpoint distance, while Traj-LeWM generally improves open-loop latent rollout accuracy over a five-step planning horizon. The improvement is favorable on three tasks and indistinguishable from zero on Cube.

  • B.3 Candidate-Selector Agreement: Selector metrics are averaged equally across queries, with 95% confidence intervals estimated from 20,000 bootstrap samples of complete queries.The queries use evaluation seeds 42–44 across all tasks.
  • B.3 Candidate-Selector Agreement: LTC-versus-endpoint top-1 agreement ranges from 8.3% to 25.0%, while top-5 intersection ranges from 12.5% to 52.5%.These values contrast with the much higher agreement between Traj-LeWM and LeWM endpoint selectors.
  • B.3 Candidate-Selector Agreement: Traj-LeWM and LeWM endpoint selectors agree on top-1 selection by 87.5%, 56.3%, 87.5%, and 83.3% for Cube, Two-Room, Reacher, and Push-T, respectively.Their corresponding top-5 intersections are 77.5%, 68.8%, 90.8%, and 81.7%.
  • B.4 Open-Loop Latent Rollout Consistency: Open-loop evaluation gives both models the same first three observations and identical recorded actions before autoregressively predicting the next ten latent states.The evaluation selects one window from each of 1,500 distinct expert episodes.
  • B.4 Open-Loop Latent Rollout Consistency: −0.00251 on Push-T, −0.00343 on Reacher, and −0.01269 on Two-Room are the five-step average paired rollout-error differences favoring Traj-LeWM.The corresponding 95% confidence intervals are [−0.00440, −0.00064], [−0.00394, −0.00291], and [−0.01535, −0.01013].
  • B.4 Open-Loop Latent Rollout Consistency: +0.00030 on Cube is the five-step average paired rollout-error difference, indicating no clear advantage for either model on that task.Its 95% confidence interval is [−0.00122, 0.00182].

C.1 Data and Paired Evaluation

Experiments used a Franka FR3 with synchronized visual, state, and action recordings, and compared LeWM with Traj-LeWM under a paired evaluation protocol. Success required avoiding collision termination and reaching a seven-joint distance of at most 0.11 rad.

  • Data: Experiments used a Franka FR3 at 30 Hz, recording eight-dimensional states, eight-dimensional actions, and three synchronized 640 × 480 RGB streams.States and actions covered seven arm joints and the gripper.
  • Paired protocol: 20 paired episode–start samples gave LeWM and Traj-LeWM the same initial observation and target joint state, with each model executed once.The target occurred 45 recorded 30-Hz steps after start, equivalent to 1.5 s, 15 10-Hz endpoints, or three 0.5-s macro steps.
  • Success criterion: A rollout succeeded only without manual collision termination and when its smallest reached seven-joint distance satisfied dmin_i,m ≤0.11 rad.The threshold was the rounded Euclidean norm of retained per-joint tolerances, whose norm was 0.108972 rad.
Loading 2608.14125v1…