LRC-JEPA: Disentangling Dynamics and Residual Context for Efficient World Models
Organizations: Independent researcher
Abstract
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 and learned-query residual-context embeddings . Only is propagated by the dynamics model and used for planning, while 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.
Figures & tables
| Model | LRC-JEPA | LWM | V-JEPA2-GAP-AC | DINO-WM | EB-JEPA |
|---|---|---|---|---|---|
| Num. of params. | 5.5M | 5.5M | 303.9M | 22.1M | 1.1M |
| Full planning time (s) | 0.45 | 0.48 | 0.83 | 57.56 | 0.20 |
| Quantity | Linear | MLP | ||||
|---|---|---|---|---|---|---|
| LWM | LRC-JEPA | LWM | LRC-JEPA | |||
| Joint pos. | 0.293 / 0.698 | 0.274 / 0.723 | 0.572 / 0.500 | 0.311 / 0.692 | 0.305 / 0.696 | 0.444 / 0.653 |
| Joint vel. | 0.969 / 0.111 | 0.734 / 0.392 | 0.990 / 0.039 | 0.999 / 0.068 | 0.893 / 0.312 | 1.013 / 0.026 |
| EE pos. | 0.019 / 0.990 | 0.011 / 0.995 | 0.380 / 0.685 | 0.009 / 0.995 | 0.007 / 0.997 | 0.144 / 0.925 |
| EE yaw | 0.989 / 0.073 | 0.930 / 0.250 | 0.876 / 0.336 | 1.009 / 0.100 | 1.027 / 0.128 | 0.897 / 0.321 |
| Model | Push-T | Bridge recall@1/5/10 | Bridge rank pct. |
|---|---|---|---|
| LRC-JEPA | 96.0 | 0.38/0.64/0.72 | 0.061 |
| w/ stop-gradient | 92.0 | 0.10/0.28/0.32 | 0.286 |
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
| Quantity | Model | Linear | MLP | ||
|---|---|---|---|---|---|
| MSE | MSE | ||||
| Agent position | LRC-JEPA | 0.066 | 0.967 | 0.022 | 0.989 |
| LWM | 0.044 | 0.978 | 0.018 | 0.991 | |
| Block position | LRC-JEPA | 0.021 | 0.989 | 0.005 | 0.997 |
| LWM | 0.021 | 0.990 | 0.006 | 0.997 | |
| Block angle | LRC-JEPA | 0.175 | 0.905 | 0.049 | 0.974 |
| Quantity | Model | Linear | MLP | ||
|---|---|---|---|---|---|
| MSE | MSE | ||||
| Agent position | LRC-JEPA | 0.002 | 0.999 | 0.000 | 1.000 |
| LWM | 0.008 | 0.996 | 0.001 | 1.000 | |
| Quantity | Linear | MLP |
|---|---|---|
| Joint position | 0.272 / 0.731 | 0.306 / 0.711 |
| Joint velocity | 0.746 / 0.397 | 1.002 / 0.088 |
| End-effector position | 0.009 / 0.996 | 0.019 / 0.991 |
| End-effector yaw | 0.891 / 0.323 | 0.923 / 0.299 |
| Gripper | 0.058 / 0.970 | 0.060 / 0.969 |
| Block position | 0.004 / 0.998 | 0.011 / 0.995 |
| Probe target | Linear | MLP | Linear | MLP |
|---|---|---|---|---|
| Gripper centroid | 0.724 / 0.690 | 0.604 / 0.735 | 1.193 / 0.321 | 0.858 / 0.402 |
| Moving-object centroid | 1.001 / 0.484 | 0.879 / 0.473 | 1.528 / 0.225 | 0.995 / 0.323 |
| Moving-object RGB | 2.440 / 0.131 | 1.141 / 0.269 | 1.914 / 0.208 | 1.179 / 0.265 |
| Static-object centroid | 1.841 / 0.252 | 1.277 / 0.259 | 1.508 / 0.246 | 1.239 / 0.267 |
| Static-object RGB | 2.324 / 0.036 | 1.208 / 0.326 | 2.035 / 0.137 | 0.999 / 0.413 |
| Probe target | Linear | MLP |
|---|---|---|
| Gripper centroid | 10.829 / 0.299 | 0.703 / 0.690 |
| Moving-object centroid | 15.260 / 0.144 | 0.888 / 0.460 |
| Moving-object RGB | 13.374 / 0.165 | 1.277 / 0.258 |
| Static-object centroid | 13.174 / | 1.219 / 0.289 |
| Static-object RGB | 11.721 / 0.204 | 1.200 / 0.347 |
| Background centroid | 4.750 / 0.156 | 0.598 / 0.571 |