PaperlayerView exact revision on arXiv

Source-linked AI summary

LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models

Luzhe Huang, Lei Chu, Jingyi Liang, Yuhuan Zhao

arXiv:2609.34375v1cs.LGcs.AIcs.RO

TL;DR

Reward-free visual goal planning must retain action-dependent state and context needed to distinguish image goals in a compact representation. LRC-JEPA routes dynamics and persistent context into separate representations, and its 5.5M active-parameter encoder improves over parameter-matched JEPA and surpasses larger pretrained competitors across simulated and real-world settings.

  • Problem

    Reward-free image-goal planning needs representations that preserve action-dependent state while distinguishing image goals without reward supervision.

    Source: §1 · p. 1

  • Method

    LRC-JEPA routes action-conditioned dynamics into a compact predictive latent z used for planning and persistent context into a residual embedding u used for reconstruction.

    Source: §3.2 · p. 4, §3.4 · p. 6

  • Results

    Across simulated and real-world settings, LRC-JEPA’s 5.5M active-parameter encoder improves over parameter-matched JEPA and surpasses up to 55× larger pretrained competitors on key planning metrics and planning time.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Theory, physical-state probes, and reconstruction interventions consistently support the intended functional separation.

    Source: §6 · p. 10

    The approach assumes nuisance context is approximately stable over the training horizon, and real-world evaluation uses offline action recovery rather than closed-loop robot execution.

    Source: §6 · p. 10

Abstract

from arXiv · show

Compact JEPA world models enable efficient latent-space planning, but low-dimensional representation trained under reward-free self-supervision must encode both action-conditioned dynamics and predictable visual context. This competition can entangle controllable state with high-rank nuisance appearance and degrade planning as scenes become more complex. We introduce LRC-JEPA, a lightweight end-to-end world model that routes information into a compact predictive latent $\mathbf{z}$ and learned-query residual-context embeddings $\mathbf{u}$. Only $\mathbf{z}$ is propagated by the dynamics model and used for planning, while $\mathbf{u}$ captures temporally persistent information for cross-attention reconstruction; a differentiable residual connection encourages the latent to retain complementary dynamic content. Under explicit assumptions, we show that the resulting representation is sufficient, minimal, nuisance-invariant, and disentangled. Across four simulated control environments, LRC-JEPA improves average planning success over a parameter-matched JEPA baseline by 9 percentage points and matches or exceeds substantially larger pretrained models. On the real-world Bridge-v2 set, its 5.5M-parameter active encoder outperforms DINO-WM (22.1M) and V-JEPA2 (303.9M) encoders while also enabling faster planning. Physical-state probes, reconstruction interventions, and ablations confirm the effectiveness of LRC-JEPA's representation disentanglement.

1 Introduction

Reward-free visual goal planning asks compact world models to preserve action-relevant dynamics alongside context needed to distinguish goals. LRC-JEPA routes these roles into separate representations, with theoretical and empirical support for the design.

1 Introduction

  • Compact world models face a capacity–compactness trade-off: low-dimensional encoders can discard action-relevant transitions or entangle controllable factors with predictable context.This is acute in reward-free image-goal planning, where future goals may depend on context that cannot be discarded a priori.Source: §1 · p. 1
  • LRC-JEPA assigns predictive dynamics to compact z and slow-varying or persistent visual context to complementary u.It is a lightweight end-to-end JEPA world model with a compact predictive latent and learned-query residual-context embeddings.Source: §1 · p. 2
  • Under explicit assumptions, LRC-JEPA satisfies sufficiency, minimality, nuisance invariance, and coordinate disentanglement, with u and z sufficient for context and dynamic state, respectively.Source: §1 · p. 2
  • Across simulated and real-world settings, LRC-JEPA’s 5.5M active-parameter encoder improves over parameter-matched JEPA and surpasses up to 55× larger pretrained competitors on key planning metrics and planning time.Probes and interventions support the intended allocation of kinematics to z and residual context to u.Source: §1 · p. 2

2 Related Work

Prior visual world models range from pixel prediction to compressed latent dynamics, while JEPA methods predict representations rather than reconstructing every input detail. LRC-JEPA separates action-conditioned prediction from reconstruction-supported context without prescribing a semantic partition.

2 Related Work

  • Visual world models have shifted from predicting pixels toward compressed latent states for planning or policy learning, with latent dynamics methods using observation, reward, or value objectives.Source: §2 · p. 2
  • JEPAs predict target representations, but preventing collapse does not inherently determine which physical variables a predictive representation preserves.DINO-WM and V-JEPA2 retain richer pretrained features, while LWM applies teacher-free distributional regularization to compact action-conditioned dynamics.Source: §2 · p. 2
  • Prior work separates appearance and dynamics, but reward-free predictive objectives can favor slowly varying distractors over changing task state.LRC-JEPA's functional bias is not a prescribed semantic partition; slowly varying quantities can still matter for control.Source: §2 · pp. 2–3
  • LRC-JEPA uses action-conditioned prediction to shape z and reconstruction with temporal invariance to encourage u to retain complementary, slowly varying context.The context stream and decoder support representation learning, while planning operates only on z.Source: §2 · p. 3

3 Method

LRC-JEPA divides information between a compact predictive latent z and residual-context embeddings u, then uses only z for prediction and planning. The method motivates this split through limitations of single-latent LWM and gives conditional theoretical guarantees for the proposed representation.

3.1 Theoretical Analysis Foundations

  • LWM’s SIGReg objective can yield a prediction-sufficient latent but does not guarantee minimality, because no training term decreases I(Zt; Xt).SIGReg’s population minimizer prevents collapse, fixes scale, and decorrelates features, but does not prevent control-irrelevant nuisance information from occupying the latent.Source: §3.1 · pp. 3–4

3.2 Model Architecture

  • LRC-JEPA assigns action-dependent state to predictive z and slowly varying context to reconstruction-only u, with the streams joined by the decoder.The shared visual encoder supplies the class token for z and learned-query pooling of the full token set for u; the context bottleneck is intended to discourage frame-wise shortcuts.Source: §3.2 · pp. 4–5

3.3 Training Objective

  • Four objectives train LRC-JEPA: one-step latent prediction, anti-collapse regularization, temporal context regularization, and patch reconstruction.SIGReg regularizes z, while a temporal VICReg-style objective encourages u to remain stable within a trajectory while varying across trajectories.Source: §3.3 · pp. 5–6

3.4 Theoretical Analysis for The Proposed Method

  • Under Assumptions 1, 4, and 5, every global minimizer has context-sufficient u, state-sufficient z conditioned on u, minimal prediction-sufficient z, nuisance invariance, and disentangled coordinates.The assumptions specify a two-factor Gaussian world with trajectory-constant appearance nuisance independent of hidden state.Source: §3.4 · p. 6
  • Residual-connected reconstruction rewards complementary dynamic information in z, whereas detaching z can leave a predictable but insufficient state.The paper’s theoretical analysis attributes this to information already represented by u not further reducing reconstruction residual under squared reconstruction loss.Source: §3.4 · p. 6

4 Results

LRC-JEPA is evaluated on simulated planning, Bridge-v2 expert-action recovery, physical-quantity probes, and reconstruction interventions. Its planning results improve on parameter-matched LWM, while probes and interventions show a division between motion-related information in z and persistent context in u.

4.1 Latent Planning Performance

  • 9 percentage points: LRC-JEPA improves average planning success over parameter-matched LWM across four simulated environments.It achieves the highest success on Push-T and Reacher and remains competitive on TwoRoom and Cube with pretrained competitors.Source: §4.1 · p. 7, Figure 3 · p. 7
  • On Bridge-v2, recall@1/5/10 is 0.38/0.64/0.72 for LRC-JEPA, versus 0.36/0.54/0.62 for V-JEPA2-GAP-AC and 0.34/0.54/0.58 for DINO-WM.LRC-JEPA leads all four offline metrics; its active encoder has 5.5M parameters, versus 303.9M for V-JEPA2-GAP-AC and 22.1M for DINO-WM.Source: §4.1 · p. 7, Figure 4 · p. 7

4.2 Physical Quantity Probing

  • LRC-JEPA outperforms LWM in 28 of 32 Cube comparisons, with joint-velocity correlations rising from 0.111 and 0.068 to 0.392 and 0.312 for linear and MLP probes.Across probes, z better predicts kinematic and control-related quantities, while u better predicts slowly varying yaw; on Bridge-v2, z better captures moving-object centroids and u generally better captures RGB values.Source: §4.2 · p. 8, Table 2 · p. 8

4.3 Reconstruction Interventions

  • Within-trajectory u shuffling causes little reconstruction degradation, while cross-trajectory u replacement changes slowly varying orientations and z shuffling changes object locations.The reported unaltered versus within-trajectory-shuffled scores are PSNR 29.58 versus 29.26 and SSIM 0.923 versus 0.919; z shuffling yields the lowest quality, PSNR 17.38 and SSIM 0.774.Source: §4.3 · pp. 8–9, Figure 5 · p. 9

5 Ablation Study

The ablations test how residual-context query count and reconstruction gradients affect reconstruction and planning-relevant representations. More queries mainly improve reconstruction, while blocking the residual gradient reduces planning performance, especially on Bridge-v2.

5 Ablation Study

  • Increasing residual-context queries from 1 to 16 primarily improves reconstruction MSE, PSNR, and SSIM; planning recall and z probes are less consistent, though gripper-centroid probing improves.The authors conclude that a small context bottleneck captures useful persistent information, while additional queries mainly improve reconstruction.Source: §5 · p. 9, Figure 6 · p. 10
  • SIGReg and next-latent prediction alone do not recover the same planning-relevant representation, particularly in visually diverse real-world settings.The authors attribute better separation of context from the predictive latent to the residual connection and reconstruction gradients.Source: §5 · p. 10
  • Blocking the residual gradient reduced Push-T success from 96% to 92% and Bridge-v2 recall@1/5/10 from 0.38/0.64/0.72 to 0.10/0.28/0.32, with rank percentile worsening from 0.061 to 0.286.Table 3 reports Push-T planning success and Bridge-v2 expert-action recall and rank percentile.Source: §5 · p. 9, Table 3 · p. 10

6 Conclusion and Limitations

LRC-JEPA routes dynamics through a predictive state and persistent visual context through reconstruction-only embeddings, with theory, probes, and interventions supporting this separation. The authors identify assumptions and evaluation boundaries, alongside statements on AI use, ethics, and reproducibility.

6 Conclusion and Limitations

  • On Bridge-v2, LRC-JEPA reports 120× acceleration and 38% control performance improvement in expert-action rank percentile versus a 4× larger DINO-WM.Theory, physical-state probes, and reconstruction interventions support the intended functional separation.Source: §6 · p. 10
  • The approach assumes nuisance context is approximately stable over the training horizon, and rapidly changing uncontrollable or irrelevant factors may be routed imperfectly.The decoder and context regularizer add training-time computation and sensitivity to bottleneck capacity and loss weights; real-world evaluation uses offline action recovery rather than closed-loop robot execution.Source: §6 · p. 10

AI use statement

  • The authors used generative AI for manuscript revision, review, reference checking and formatting, code and documentation editing, and feedback on methodology and result interpretation.They state that authors independently verified AI-assisted materials and take responsibility for the final manuscript and released artifacts.Source: AI use statement · p. 11

Ethics statement

  • Bridge-v2 evaluation measures offline action ranking and does not establish the safety or reliability of deployment on physical robots.Physical-system use requires additional validation and appropriate operational safeguards; datasets, software, and pretrained models are subject to their licenses and access conditions.Source: Ethics statement · p. 11

Reproducibility statement

A Implementation Details

The implementation details specify the planning-time model, baseline construction, capacity and nuisance interventions, and evaluation procedures for probes and speed. The Bridge-v2 instantiation includes a larger training-time decoder, while planning loads only the predictive-state path.

A.1 Neural Networks

  • The Bridge-v2 model uses a ViT-Tiny encoder and a 6-layer predictor; its approximately 21.1M total parameters include the decoder, while planning loads 5.5M active encoder parameters.The compact V-JEPA2-GAP-AC baseline uses a frozen pretrained ViT-L/16 encoder, globally averaged 1024-dimensional states, and a six-layer action-conditioned predictor.Source: §A.1 · p. 15

A.2 Details on Capacity–complexity Stress Tests for Compact JEPA planning

  • The PCA intervention projects encoded and predicted latents into a shared leading-k subspace, restricting effective rank while preserving the 192-dimensional tensor shape.It uses k values from 32 to 192; k = 192 is a no-op, so the intervention measures informational capacity rather than native model computation.Source: §A.2 · p. 15
  • The static-pattern intervention blends a trajectory-persistent 8 × 8 binary grayscale pattern into current and goal observations at opacity α values from 0 to 0.50.Patterns vary across trajectories but remain fixed over time and independent of actions; the same 50 trajectories and planning seeds are used across opacity conditions.Source: §A.2 · p. 15

A.3 Details on Physical Quantity Probes

  • Physical-quantity probes use frozen frame representations, with Bridge-v2 object labels from segmentation masks and an 80/20 split by episode to prevent trajectory leakage.The probes compare z, mean-pooled u, and their concatenation using ridge regression and a two-layer MLP.Source: §A.3 · pp. 15–16
  • Planning speed is measured on one NVIDIA GeForce RTX 5090 GPU using matched CEM settings, with full planning time averaged over 50 trajectories.Timing includes online current- and goal-image encoding but excludes checkpoint loading, environment stepping, video generation, and result serialization.Source: §A.3 · p. 16

A.4 Implementation Details

This material describes LRC-JEPA’s training and planning choices, contrasts its division of representational roles with prior designs, states the assumptions used in its theory, and explains what SIGReg’s decomposition does and does not constrain.

A.4 Implementation Details

  • SIGReg regularizes latent distributions to prevent collapse, while temporal VICReg stabilizes context within trajectories and preserves variation across them.SIGReg uses random one-dimensional projections; temporal regularization computes variance and covariance across clip-level means.Source: §A.4 · pp. 16–17
  • At a global optimum, LRC-JEPA’s predictive stream is compact and planned through, while its residual context supports reconstruction rather than dynamics prediction.Planning rolls the predictor forward in z and scores terminal predicted states against encoded goals; CEM candidates are replanned in closed loop.Source: §A.4 · p. 17, §A.5 · p. 18

A.5 Relation to Prior Designs

  • Unlike LWM’s single predictive latent and MC-JEPA’s shared motion-content representation, LRC-JEPA assigns prediction to z and temporally regularized residual context to u.Its decoder queries the compact [zt; udec] representation rather than all visible patch tokens.Source: §A.5 · pp. 17–18

B.1 Definitions and assumptions

  • The theoretical analysis assumes a noisy surrogate encoder and a single-factor Gaussian controlled world with spectral separation, action variation, and population optimality.The single-factor identification result also assumes the predictive latent dimension matches the state dimension and that the encoder and predictor attain a global minimum.Source: §B.1 · pp. 18–19

B.2 The decomposition lemma

  • The decomposition lemma splits a factorized-prior KL into index-code mutual information, total correlation, and coordinate-wise matching, which can be controlled separately.Coordinate-wise matching can vanish while index-code mutual information takes compatible values and total correlation remains positive.Source: §B.2 · p. 19
  • SIGReg’s Gaussian aggregate target makes total correlation and coordinate-wise mismatch vanish together, but leaves index-code mutual information unconstrained.Thus, matching the aggregate distribution does not by itself constrain the information rate.Source: §B.2 · p. 19

B.3 Proof of Proposition 1(Sufficiency)

The proof uses SIGReg’s Gaussian aggregate target to establish an orthogonal relationship between the learned predictive latent and the underlying state.

B.3 Proof of Proposition 1(Sufficiency)

  • Under Assumptions 2 and 3, zero SIGReg implies a standard-Gaussian aggregate, and the deterministic predictive representation equals the state up to an orthogonal transformation.The proof then recovers the state by applying the transpose of that orthogonal transformation, establishing state sufficiency.Source: §B.3 · pp. 19–20

B.4 Proof of Proposition 1(Minimality)

The proof argues that the single-latent LWM objective does not enforce minimality, even though prediction loss encourages retention of transition-relevant information.

B.4 Proof of Proposition 1(Minimality)

  • SIGReg depends only on the aggregate latent distribution, so the LWM objective neither explicitly targets nor enforces minimality.The passage notes that SIGReg’s Gaussian target primarily prevents collapse, while prediction loss encourages retention of transition-relevant information.Source: §B.4 · p. 20

C.1 Optimality of LRC-JEPA

Under a two-factor Gaussian-world model and stated capacity and optimum assumptions, the proof separates trajectory-constant appearance context from predictive state. It then establishes context and state sufficiency, minimality, nuisance invariance, and coordinate disentanglement for the predictive latent.

C.1 Optimality of LRC-JEPA

  • Under Assumptions 1, 4, and 5, every global minimizer makes u sufficient for appearance and z sufficient for state conditioned on u, while z is minimal, nuisance-invariant, and coordinate-disentangled.The theorem assumes trajectory-constant appearance independent of state, minimum representation rate, and a context embedding constant within each trajectory while retaining cross-trajectory variance.Source: §C.1 · p. 20
  • At the optimum, SIGReg makes the predictive latent standard Gaussian with zero total correlation, and temporal regularization makes context constant within a trajectory and equal to a function of appearance.These proof steps establish coordinate disentanglement and restrict the context stream to shared trajectory context.Source: §C.1 · p. 20
  • Exact reconstruction and context sufficiency imply that z retains state information conditioned on u.The proof uses injectivity of the observation map to recover state from the reconstructed observation once appearance is fixed.Source: §C.1 · p. 21
  • The proof derives nuisance invariance from the routing constraint and independence assumptions, then establishes minimality by showing z attains the lower bound for prediction-sufficient representations.The predictable signal is assumed to span the state.Source: §C.1 · p. 21
  • Under the assumed linear-Gaussian controlled dynamics, z minimizes next-latent mean squared error and is equivalent to state up to an orthogonal transformation.The passage attributes this equivalence to Klindt et al. (2026) and Zhang et al. (2026).Source: §C.1 · p. 21

D Additional Results

Additional results probe physical information in LRC-JEPA’s representations across simulated and real-world settings, and visualize reconstructions under temporal and cross-trajectory context interventions.

D Additional Results

  • Push-T and TwoRoom probes compare LRC-JEPA and LWM predictive latents using ridge-regression and MLP heads on simulator-provided physical quantities.The probes use frozen, converged encoder outputs and held-out data; Table 4 and Table 5 report MSE and Pearson correlation.Source: D Additional Results · p. 21, Table 4 · p. 22, Table 5 · p. 22
  • Cube probes evaluate physical quantities from LRC-JEPA’s concatenated z ⊕ u representation using MSE and Pearson correlation.The table caption defines ⊕ as concatenation.Source: D Additional Results · p. 21, Table 6 · p. 22
  • Bridge-v2 probes assess object-centric physical quantities using z, mean-pooled Q = 2 context embeddings, and their concatenation.One valid moving, static, and background object is selected per frame; Table 7 marks better results by probe type and metric.Source: D Additional Results · p. 21, Table 7 · p. 23, Table 8 · p. 23

D.1 Reconstruction Visualization

  • Push-T and TwoRoom reconstructions are shown under within-trajectory frame shuffling and context exchanged across trajectories.Figure 7 covers Push-T; Figure 8 shows the same interventions for TwoRoom.Source: Figure 7 · p. 24, Figure 8 · p. 25

Related papers

Loading 2609.34375v1…