PaperlayerView exact revision on arXiv

Source-linked AI summary

XP-JEPA: Cross-Predictive Physics Grounding for Forecastable Latent Dynamics

Kehan Wen, Ziming Li, Siyuan Luo, Fan Shi

arXiv:2608.24044v1cs.LG

TL;DR

Self-predictive latent world models can learn transitions that are predictable yet weakly constrained by physical evolution, limiting rollout-based control. XP-JEPA uses privileged physical trajectories to ground visual latent dynamics through shared cross-predictive learning, then discards the physical branch; across six evaluation subfamilies, it reduces rollout drift and improves mean control success.

  • Problem

    Self-predictive models may co-adapt representations and predictors to physically weakly constrained transitions, while state decodability does not ensure forecastable dynamics or planning utility.

  • Method

    XP-JEPA separately encodes visual observations and physical states, advances both histories with a shared action-conditioned predictor, and matches each prediction to both future representations before removing the physical branch.

  • Results

    Across the multi-task suite, XP-JEPA reduces relative rollout drift from 0.361 to 0.104 and raises mean control success from 53.6% to 78.2%.

  • Takeaways & Limitations

    Cross-predictive physical grounding can improve forecastable latent dynamics for rollout-based control without privileged inputs at test time.

  • Takeaways & Limitations

    Results are limited to simulation, paired privileged trajectories during training, and a fixed rigid-body state schema.

Abstract

from arXiv · show

Latent world models plan by predicting how candidate actions transform learned representations. In self-predictive models, however, the encoder and predictor are optimized jointly and can co-adapt to latent transitions that are easy to predict but only weakly constrained by the physical evolution of the scene. We introduce the cross-predictive JEPA (XP-JEPA), which grounds visual latent dynamics in privileged physical trajectories. XP-JEPA separately encodes visual observations and physical states, advances both through a shared action-conditioned predictor, and matches each prediction to both future representations. This objective encourages unified latent dynamics across the two modalities, grounded in the underlying physical transitions. The physical branch is discarded after training, leaving a visual-only model at deployment. On a multi-task suite spanning six evaluation subfamilies, XP-JEPA reduces rollout drift of a newly fitted predictor from $0.361$ to $0.104$ and increases mean control success from $53.6\%$ to $78.2\%$. Direct physical-state regression raises position decodability but leaves forecastability and control near the visual-only baseline. Cross-predictive physical grounding can therefore produce more forecastable latent dynamics for rollout-based control without privileged inputs at test time.

1 INTRODUCTION

XP-JEPA addresses a gap in latent world models: jointly learned representations can support predictable but physically weakly constrained transitions. It grounds visual predictive dynamics in privileged physical trajectories, improving forecastability and control while separating snapshot decodability from transition quality.

  • Motivation: Self-predictive encoders and predictors may co-adapt to latent transitions that are easy to predict but weakly constrained by physical evolution.This creates ambiguity when visually similar states evolve differently under the same action.
  • Motivation: Privileged physical trajectories provide an external reference for constraining both preserved information and its action-conditioned evolution.The goal is to constrain predictive evolution rather than simply encode more physical information.
  • Evaluation framing: Decodability, forecastability, and planning utility are distinct: accessible physical variables do not necessarily yield predictable evolution or action-relevant latent geometry.The experiments explicitly separate these properties.
  • Results: REGRESS reaches 53.7% versus VISUAL’s 53.6% suite success despite higher object-position decodability, showing that position probes do not capture gripper–object relations.This contrast supports distinguishing snapshot content from transition structure.
  • Approach: XP-JEPA advances visual and physical histories with a shared action-conditioned predictor and matches predictions to both modalities’ corresponding futures.The physical branch is discarded after training, leaving the visual-only deployment architecture and planner.
  • Results: 53.6% to 78.2%: XP-JEPA increases mean multi-task control success across the evaluation suite.The result is reported alongside improved forecastable visual dynamics across the suite and all four matched single-task environments.

2 RELATED WORK

Prior work uses privileged information for representation transfer, behavior supervision, and physical grounding, but the literature leaves open how to supervise predictive latent transitions across heterogeneous tasks.

  • Predictive latent world models: Predictive latent world models include reconstruction-based, value-centric, and representation-predictive approaches that differ in the structure their representations preserve.Representation-predictive methods model future features directly for planning.
  • Privileged supervision: Privileged-information methods transfer additional training-time signals through distillation, cross-modal alignment, state regression, privileged sensing, or asymmetric actor–critic training.These signals are available during training but unavailable at deployment.
  • Physical grounding: Recent approaches bring privileged physical information closer to learned dynamics through state-to-image transfer, privileged sensors, representation alignment, or privileged prediction.XP-JEPA is positioned among methods that ground predictive dynamics using physical information.

3 CROSS-PREDICTIVE JEPA

XP-JEPA adds a physical predictive view to a visual JEPA, using shared action-conditioned prediction and corresponding cross-modal targets to ground latent evolution. A unified rigid-body interface supports heterogeneous scenes, while only the visual branch remains at deployment.

  • Core objective: XP-JEPA treats privileged physical state as a second predictive view of the same scene trajectory.The physical encoder maps instantaneous state to a physical latent, while both modalities infer motion from temporal context and action.
  • Core objective: The baseline visual model jointly learns an observation encoder and action-conditioned predictor to estimate future observation representations.XP-JEPA retains visual self-prediction and adds a physical branch with physical self-prediction and cross-modal terms.
  • Predictor sharing: Visual and physical histories are advanced by the same action-conditioned predictor without an explicit modality indicator.The encoders remain jointly optimized, with branch-specific isotropy regularization controlling collapse.
  • Corresponding cross-prediction: Four source–target terms combine within-modality self-prediction with cross-modal prediction from either history to the other modality’s future.The predicted visual transition is constrained by the corresponding physical future, not only the visual future.
  • Unified privileged-state interface: A unified rigid-body schema represents object geometry and pose, end-effector state, and a learned table token across heterogeneous configurations.Objects are represented by geometry rather than identity, avoiding category labels and task-specific state fields; unused slots are masked.
  • Deployment: At deployment, the physical encoder is discarded, while the visual encoder, predictor, inference architecture, and computational cost match the visual-only baseline.Planning rolls candidate action sequences through latent dynamics and scores predicted terminal latents against an encoded goal.

4 EXPERIMENTS

Across multi-task and matched single-task evaluations, XP-JEPA improves forecastability and, on the multi-task suite, control under matched planning conditions. Ablations indicate that cross-modal correspondence is central to control gains, while direct cross-modal prediction further reduces rollout drift.

  • Evaluation setup: XP-JEPA evaluates forecastability by freezing each visual encoder, discarding its co-trained predictor, and fitting the same action-conditioned predictor from scratch.This separates representation forecastability from encoder–predictor co-adaptation during training.
  • Forecastability: Relative rollout drift falls from 0.361 for VISUAL to 0.104 for XP-JEPA on the multi-task suite.The separation persists at every rollout horizon, including beyond h = 3.
  • Forecastability: Forecastability improves across all four matched single-task environments, with relative drift decreasing on Two-Room, OGBench-Block, Push-T, and Reacher.The reported decreases are 0.503 to 0.221, 0.507 to 0.269, 0.399 to 0.258, and 0.313 to 0.233, respectively.
  • Control: Mean multi-task task success rises from 53.6% to 78.2%, a gain of 24.6 percentage points under the same planner, action prior, and candidate budget.The improvement holds across all six interaction families.
  • Decodability versus control: State regression raises manipulated-object position decodability to R2 = 0.991 but leaves control at 53.7% and relative rollout drift at 0.373.This separates snapshot decodability from forecastable transition structure; on single-task environments, control gains are task dependent.
  • Ablations: CROSS-ONLY reaches 73.0% control with comparable rollout drift to XP-JEPA, whereas SHARE-ONLY reaches 56.7%, close to VISUAL.The ablations indicate that corresponding cross-prediction contributes more to control than predictor sharing alone.

5 DISCUSSION AND CONCLUSION

XP-JEPA grounds visual latent evolution in privileged physical trajectories, improving forecastability and mean control across the multi-task suite. Its benefits depend on correct correspondence and do not guarantee higher success in every environment.

  • 0.104 relative rollout drift and 78.2% mean control success are achieved by XP-JEPA, versus 0.361 and 53.6% for the visual-only baseline.
  • Direct state regression improves position decodability but leaves forecastability and control near the visual-only baseline.
  • Correct visual–physical correspondence supports control, while direct cross-modal prediction most clearly improves forecastability beyond alignment.
  • Forecastability alone is insufficient because predicted trajectories must preserve distinctions relevant to action selection.
  • XP-JEPA retains lower drift after replacing the co-trained predictor, but improved forecastability does not guarantee higher success in every environment.
  • The results are limited to simulation, paired privileged trajectories during training, and a fixed rigid-body state schema.Extending the approach to real-world observations and broader physical interactions remains future work.

A SUITE AND EVALUATION DETAILS

The multi-task suite evaluates 22 object–task configurations across six interaction families using held-out, consistently sampled scenarios. Evaluation uses executed physical trajectories and unchanged benchmark procedures where specified.

  • 22 object–task configurations over 13 assets form the multi-task corpus, with 450 clean episodes per configuration before noise augmentation.A single model is trained jointly across all configurations.
  • Success is determined from executed physical trajectories, using asset-specific positional tolerances and symmetry-aware orientation evaluation.
  • 300 held-out scenarios per model are sampled across configuration–subtask pairs from 266 source episodes using a fixed seed.Segments from the same source episode are clustered for statistical testing.
  • Six evaluated interaction families are annotated subtasks within episodes, with scenarios keyed by configuration–subtask pairs.Quotas are filled round-robin across configurations containing each subtask.
  • The four LeWM environments use the released success predicates and evaluation loop unchanged, with N = 150 paired scenarios per controller.

B FULL CONTROL RESULTS

XP-JEPA consistently outperforms VISUAL across the reported control evaluations, with statistically supported paired differences and higher mean suite control than CROSS-ONLY.

  • All six seed-level XP-JEPA comparisons against VISUAL and DISTILL yield pMC ≤10^-5.The paired differences against VISUAL are +28.3, +22.0, and +23.7 percentage points across seeds.
  • Table 6 reports per-family control success percentages as mean ± sample SD over three training seeds, ordered by the XP-JEPA–VISUAL gap.
  • XP-JEPA exceeds VISUAL in every seed on both Two-Room and OGBench-Block, with the largest p-value equal to 0.023.
  • 79.0/76.7/79.0 mean suite control for XP-JEPA exceeds 72.3/74.0/72.7 for CROSS-ONLY.The training run is the unit of analysis for this comparison.

C FORECASTABILITY PROTOCOL AND ADDITIONAL RESULTS

Forecastability is measured with independently fitted predictors on frozen representations and matched held-out episodes. XP-JEPA retains a forecastability advantage under physical decoding, fresh predictors, fixed-action rollouts, and rank-controlled projections.

  • Forecastability protocol: Fresh-predictor probes freeze each encoder, discard the original predictor, and train equal-capacity predictors from scratch to isolate representation forecastability.
  • Forecastability protocol: The action-conditioned probe uses the three most recent latent frames and action chunks to predict the next latent as a residual update.
  • Forecastability protocol: Autoregressive 20-step evaluation measures prediction under actions, while normalization by a constant-latent baseline removes credit for temporal persistence.
  • Decodability: The manipulated-object decoding probe maps one frozen latent frame to a 15-dimensional target including position, orientation, and end-effector position.
  • Multi-task results: XP-JEPA has lower raw rollout error and normalized drift than VISUAL and state-regression baselines, while REGRESS and DISTILL provide the highest position decodability.
  • Ablation interpretation: SHUFFLE has lower normalized drift partly because its Copy denominator is 15.11, yet its raw rollout error is 2.09 versus VISUAL’s 4.43 and control fails.This supports the distinction between forecastability and control under incorrect visual–physical supervision.
  • Trajectory visualization: 1.4 cm mean terminal error for XP-JEPA compares with 3.1 cm for VISUAL under identical stacking actions and starting observations.The decoded object position is compared across three planning cycles.
  • Rank-controlled analysis: At k = 57, XP-JEPA and VISUAL have measured ranks of 51.7 and 52.8, but drift remains 0.121 versus 0.453.The forecastability gap persists after matching effective rank.

D PRIVILEGED-STATE INTERFACE

The privileged-state interface represents physical scenes with structured object, effector, and table tokens. Inputs describe instantaneous configurations, while temporal histories provide motion context.

  • State representation: Each object token combines canonical geometry, bounding-box half-extents, rotation, and position.The geometry descriptor is an 83-dimensional average-pooled signed-distance field.
  • State representation: The effector token contains rotation, position, and finger opening, while a learned table token completes the token set.
  • Input scope: The privileged input excludes velocities, wrenches, contact forces, and goal-relative quantities.Both branches receive length-H histories and infer motion from temporal context and action.

E IMPLEMENTATION DETAILS

The implementation uses matched visual and physical encoders with a shared action-conditioned predictor, multi-task noisy data, and latent-space CEM planning. Training preserves physically consistent perturbations and discards privileged inputs at deployment.

  • Architecture: The visual encoder is a ViT-Tiny, while the physical encoder is a Transformer consuming the unified privileged-state token set.The visual encoder uses 192-dimensional features; the physical encoder has width 192 and hidden width 256.
  • Architecture: The predictor advances H = 3 latent frames with action-conditioned Transformer dynamics, and four source–target terms receive equal weight.Actions enter through zero-initialized AdaLN, with five-step action chunks and a one-step prediction offset.
  • Objective: XP-JEPA uses online encoder outputs without stop-gradient, an exponential-moving-average encoder, or a separate target network.
  • Optimization: Multi-task models train for 40 epochs using AdamW, bfloat16 precision, gradient clipping, and a warmup cosine-annealing schedule.No image augmentation is applied; diversity comes from corpus-level action-noise augmentation.
  • Planning: Planning searches five macro-steps with 300 candidates, 30 CEM iterations, and 30 elites before executing 25 environment actions.Latent action-space candidates decode to temporally coherent action chunks.
  • Data: The multi-task corpus contains 29,260 episodes and approximately 1.44M transitions across 22 configurations.It combines one clean shard with four action-noise shards.
  • Data collection: Noise is injected during collection at waypoint and action levels, producing physically consistent rollouts of perturbed policies.Waypoint perturbations persist across each replan segment, while action noise follows an Ornstein–Uhlenbeck process.
  • Data collection: Noise tiers are calibrated by family-specific factors to target failure rates from 20–30% through at least 98%.Yaw waypoint noise is disabled throughout.

F ABLATION DETAILS

The ablations isolate correspondence, predictor sharing, state regression, distillation, and explicit alignment while keeping core training conditions matched. Their designs distinguish transition grounding from snapshot supervision and shared dynamics.

  • Ablation controls: All ablations retain XP-JEPA’s training data, visual backbone, latent dimensionality, optimization settings, and schedule unless the intervention requires otherwise.
  • Predictor sharing: CROSS-ONLY preserves corresponding cross-prediction but uses separate visual and physical predictors.This isolates corresponding cross-prediction without a common transition model.
  • Predictor sharing: SHARE-ONLY retains the shared predictor but removes cross-modal terms, so each modality predicts only its own future representation.This isolates predictor sharing without direct cross-modal grounding.
  • Snapshot supervision: REGRESS adds a per-frame state-regression head to the unchanged visual branch and self-prediction term.The head maps visual latents to an 85-dimensional pose-based target and is unused at deployment.
  • Snapshot supervision: DISTILL replaces predictive cross-modal coupling with pointwise regression while preserving sample-level physical correspondence.Each branch uses its own predictor and output projection, supervising representation content rather than corresponding predictive transitions.
  • Alignment: ALIGN-ONLY keeps shared prediction and within-modality terms but replaces cross-modal loss terms with explicit symmetric alignment.Alignment is applied to the full encoded window with λalign = 1.
  • Correspondence: SHUFFLE preserves intact physical trajectories but pairs each visual episode with a fixed incorrect partner trajectory.Pixels and actions remain on the true episode, isolating correspondence between modalities.

G QUALITATIVE ROLLOUTS

The qualitative rollout comparisons pair XP-JEPA and VISUAL from identical initial states with the same planner across six interaction families. The displayed disagreement cases show XP-JEPA succeeding where VISUAL fails.

  • Comparison protocol: Figures 8 and 9 compare both methods from the same initial state using the same planner across all six evaluated interaction families.Each example is a representative disagreement scenario, with frames cropped to the working volume.
  • Grasping, insertion, placement: In grasping, insertion, and placement, each paired scenario shows XP-JEPA completing the shared goal while VISUAL does not.XP-JEPA appears above VISUAL in each pair.
  • Pushing, stacking, tossing: In stacking, VISUAL leaves the eraser beside the block, while in tossing it leaves the object outside the carton.
Loading 2608.24044v1…