DepState-JEPA: Learning Depth-Grounded State Representation for Visual Perception and Manipulation
Organizations: Fudan University · Shanghai Innovation Institute (SII) · Shanghai Jiao Tong University · Shandong University ·
Abstract
Predictive representations for robot manipulation should connect scene geometry with the state changes induced by actions. Our intuition is that understanding what an action will change starts with understanding where things are: complementary RGB and depth observations ground a common scene representation from which to infer those changes. Guided by this intuition, we introduce DepState-JEPA, a dual-stream RGB--Depth joint-embedding predictive architecture that connects spatial understanding with action-conditioned prediction through a shared latent state, DepState. Built on V-JEPA~2.1, the framework uses depth both as a complementary observation stream and to ground both modalities in metric 3D space, aggregating shared information into DepState while retaining modality-private features for decoding. Ray-conditioned readouts predict local features without target depth; after an action-conditioned transition, the same decoders predict future RGB and depth features using the predicted state and fixed initial private features. Robot-state readouts provide recursive proprioceptive inputs; object-trajectory supervision and calibrated forward kinematics shape state learning through motion targets and joint--TCP consistency. Comparisons show stronger RGB-encoder transfer after RGB-D adaptation, feature completion, conditioned visual forecasting, and manipulation success. Within-model ablations identify contributions from depth content, metric coordinates, and future-state matching, with additional control gains from FK regularization. The project page is available at https://shawnleung-sii.github.io/depstate-jepa/.
Figures & tables
| MDE on NYUv2 | SS on VOC12 | SC on LB-Depth | ||||||
|---|---|---|---|---|---|---|---|---|
| Method | AbsRel | RMSE | mIoU | mAcc | L1 | Cos. | ||
| LeJEPA ( Balestriero & LeCun, 2025 ) | 0.137 | 0.490 | 0.059 | 85.7 | 76.0 | 85.0 | 0.352 | 0.891 |
| V-JEPA 2 ( Assran et al., 2025 ) | 0.188 | 0.642 | 0.079 | 77.5 | 63.9 | 76.8 | 0.428 | 0.843 |
| V-JEPA 2.1 ( Mur-Labadia et al., 2026 ) | 0.098 | 0.350 | 0.042 | 92.5 | 84.7 | 91.8 | 0.298 | 0.931 |
| 3D-JEPA † ( Arnaud et al., 2025 ) | Not evaluated: point-cloud input | 0.283 | 0.938 | |||||
| 1 step | 8 steps | 16 steps | 24 steps | 32 steps | ||||||
|---|---|---|---|---|---|---|---|---|---|---|
| Method | L1 | Cos. | L1 | Cos. | L1 | Cos. | L1 | Cos. | L1 | Cos. |
| V-JEPA 2-AC ( Assran et al., 2025 ) | 0.242 | 0.916 | 0.421 | 0.792 | 0.523 | 0.701 | 0.580 | 0.646 | 0.658 | 0.598 |
| V-JEPA 2.1-AC ( Mur-Labadia et al., 2026 ) | 0.171 | 0.967 | 0.342 | 0.874 | 0.476 | 0.781 | 0.554 | 0.702 | 0.614 | 0.636 |
| 3D-JEPA-AC † ( Arnaud et al., 2025 ) | 0.169 | 0.968 | 0.336 | 0.879 | 0.467 | 0.789 | 0.546 | 0.711 | 0.602 | 0.647 |
| LeMuMo-JEPA-AC † ( Cornelissen et al., 2026 ) | 0.167 | 0.969 | 0.331 | 0.882 | 0.479 | 0.778 | 0.563 | 0.694 | 0.624 | 0.627 |
| DepState-JEPA-AC (Ours) | 0.145 | 0.976 | 0.307 | 0.903 | 0.394 | 0.851 | 0.451 | 0.811 | 0.491 | 0.779 |
| Online PL | Latent MPC | Sim2Real MPC | ||||
|---|---|---|---|---|---|---|
| Method | DMC | MW | DMC | MW | Task 1 | Task 2 |
| V-JEPA 2-AC ( Assran et al., 2025 ) | 682.4 14.7 | 52.8 2.8 | 736.5 15.3 | 66.4 2.4 | 45.6 5.1 | 41.1 5.1 |
| V-JEPA 2.1-AC ( Mur-Labadia et al., 2026 ) | 763.8 15.4 | 68.6 2.3 | 807.2 15.0 | 81.2 2.6 | 66.7 6.7 | 64.4 5.1 |
| 3D-JEPA-AC † ( Arnaud et al., 2025 ) | 784.4 12.6 | 72.3 2.6 | 826.4 18.2 | 85.6 3.2 | 73.3 5.8 | 71.1 5.1 |
| LeMuMo-JEPA-AC † ( Cornelissen et al., 2026 ) | 790.4 16.1 | 73.2 2.5 | 837.5 13.6 | 87.8 2.2 | 75.6 6.9 | 74.4 8.4 |
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
| Symbol | Definition |
|---|---|
| , | Complete observation and its visible context in the JEPA objective. |
| ; | Registered RGB, metric depth, and camera intrinsics; . |
| Visible and masked token-index sets. | |
| , | Geometrically valid visible correspondences and masked targets. |
| Token, modality, encoder level, control-time, and object/part-slot indices. | |
| Context encoder, predictor, and momentum target encoder. |
| Symbol | Definition |
|---|---|
| Homogeneous coordinates of patch center . | |
| , | Metric center depth and back-projected point . |
| ; | Fused feature before/after metric encoding; is the grounded token set. |
| Learned projection of Fourier-encoded metric coordinates. | |
| ; | Shared/private projections and their output features for modality . |
| RGB/depth private features: and . |
| Symbol | Definition |
|---|---|
| Action-conditioned temporal Transformer. | |
| Joint command, configuration, instantaneous velocity, and control interval. | |
| , , | Predicted endpoint configuration, velocity, and TCP pose. |
| Independent joint-configuration and TCP-pose heads. | |
| , | Predicted object/part body increment and its readout head. |
| Object/part integration origin. |
| Symbol | Definition |
|---|---|
| Dense multi-level RGB/depth latent prediction on masked and visible targets. | |
| Masked shared-feature completion. | |
| Grouped spatial prediction: dense features plus shared completion. | |
| , , | Shared alignment, shared–private decorrelation, and variance floor. |
| Grouped factorization objective. | |
| , | Future feature and future state matching components. |
| Component | Configuration |
|---|---|
| Multi-level fusion | Per modality: , output LN. |
| Shared / private heads | Independent per modality: / , output LN. |
| Metric PE | Fourier XYZ: ; shared across modalities; final linear layer initialized to zero. |
| State aggregator | 64 learned queries, width 384, 2 blocks; query self-attention, cross-attention, and FFN. |
| Ray-query embedding | Fourier ray plus time: . |
| Spatial decoder | 6 blocks, width 384; query self-attention, cross-attention to , and FFN. |
| Setting | RGB-D Adaptation | Robot Post-training |
|---|---|---|
| Initialization | Pretrained encoders; new state and decoding modules | Adapted online/teacher; new temporal predictor and physical heads |
| Total updates | 35,000 | 20,000; every fourth update is replay |
| Resolution | 256 for 30,000 updates; 384 for the last 5000 | 384 for robot data and replay |
| Temporal input | Images / 16-frame videos at 4 fps | Single image per robot state; 16-frame replay videos |
| Effective / global micro-batch | 128 / 16 | 128 / 16 |
| Peak encoder LR |
| Online PL | Latent MPC | ||||
| Method | Seed | DMC | MW | DMC | MW |
| V-JEPA 2-AC | 0 | 671.8 | 50.8 | 719.4 | 64.2 |
| V-JEPA 2.1-AC | 0 | 757.6 | 66.4 | 799.1 | 78.8 |
| 3D-JEPA-AC † | 0 | 781.7 | 72.4 | 822.8 | 85.6 |
| LeMuMo-JEPA-AC † | 0 | 787.3 | 72.2 | 833.5 | 87.4 |
| Task 1: Drawer | Task 2: Cup | |||||
| Method | Pull | Close | Full | Pick | Place | Full |
| V-JEPA 2-AC | 66.7 3.3 | 81.1 1.9 | 45.6 5.1 | 63.3 3.3 | 73.3 3.3 | 41.1 5.1 |
| V-JEPA 2.1-AC | 86.7 3.3 | 90.0 3.3 | 66.7 6.7 | 83.3 3.3 | 87.8 1.9 | 64.4 5.1 |
| 3D-JEPA-AC † | 90.0 3.3 | 91.1 1.9 | 73.3 5.8 | 87.8 5.1 | 88.9 1.9 | 71.1 5.1 |
| LeMuMo-JEPA-AC † | 90.0 3.3 | 91.1 1.9 | 75.6 6.9 | 88.9 5.1 | 90.0 3.3 | 74.4 8.4 |
| Task 1: Drawer | Task 2: Cup | ||||||
| Method | Seed | Pull | Close | Full | Pick | Place | Full |
| V-JEPA 2-AC | 0 | 20/30 | 24/30 | 14/30 | 18/30 | 21/30 | 11/30 |
| V-JEPA 2.1-AC | 0 | 26/30 | 27/30 | 20/30 | 24/30 | 26/30 | 18/30 |
| 3D-JEPA-AC † | 0 | 27/30 | 28/30 | 23/30 | 26/30 | 27/30 | 21/30 |
| LeMuMo-JEPA-AC † | 0 | 27/30 | 27/30 | 22/30 | 27/30 | 27/30 | 22/30 |
| ID | Experimental Setting | L1 | Cos | SR |
|---|---|---|---|---|
| Ex.1 | Replace online depth content with RGB and disable metric XYZ positional encoding. | 0.408 0.021 | 0.832 0.017 | 76.8 2.9 |
| Ex.2 | Retain depth content; replace metric XYZ positional encoding with an unknown-geometry embedding. | 0.369 0.017 | 0.860 0.014 | 80.1 2.7 |
| Ex.3 | Retain metric XYZ positional encoding; replace online depth content with RGB. | 0.359 0.017 | 0.869 0.014 | 81.3 2.7 |
| Ex.16 | Full model: retain both depth content and metric XYZ positional encoding. | 0.340 0.013 | 0.880 0.011 | 83.4 2.9 |
| ID | Experimental Setting | L1 | Cos | SR |
|---|---|---|---|---|
| Ex.4 | Disable the private cross-attention residual in the modality decoders; retain private projections and factorization losses. | 0.390 0.019 | 0.852 0.016 | 83.2 2.7 |
| Ex.5 | Add unit-weight alignment of corresponding private features to the existing shared-feature alignment. | 0.376 0.018 | 0.855 0.016 | 81.5 2.9 |
| Ex.6 | Replace projected aggregation keys/values by repeated valid-token means; retain learned queries and state size. | 0.440 0.029 | 0.799 0.025 | 77.0 1.9 |
| Ex.7 | Replace unit-ray query inputs with normalized pixel vectors ; retain the query MLP. | 0.340 0.013 | 0.880 0.011 | 83.5 2.7 |
| Ex.8 | Remove completion supervision throughout training/adaptation; retain dense JEPA and temporal objectives. | 0.361 0.018 | 0.870 0.014 | 82.8 2.9 |
| Ex.16 | Full model: private decoding, shared-only alignment, learned aggregation, ray queries, and completion supervision. | 0.340 0.013 | 0.880 0.011 | 83.4 2.9 |
| ID | Experimental Setting | L1 | Cos | SR |
|---|---|---|---|---|
| Ex.9 | Remove future-state matching; retain future-feature matching and physical supervision. | 0.397 0.024 | 0.827 0.021 | 79.5 2.8 |
| Ex.10 | Remove future-feature matching; retain future-state matching and physical supervision. | 0.392 0.019 | 0.833 0.017 | 82.2 2.8 |
| Ex.11 | Remove object-trajectory supervision; retain the object head and integration interface. | 0.344 0.014 | 0.877 0.012 | 83.2 2.9 |
| Ex.12 | Remove FK consistency regularization; retain independent joint/TCP heads and their direct labels. | 0.359 0.018 | 0.864 0.016 | 82.5 2.8 |
| Ex.13 | Apply FK gradients to joint/TCP heads only via a detached state; keep direct supervision undetached. | 0.350 0.016 | 0.871 0.014 | 83.1 2.9 |
| Ex.14 | Recursively update private memory from predicted states without future observations; add a memory updater. | 0.351 0.017 | 0.871 0.015 | 83.1 2.9 |
| Seed 0 | Seed 1 | Seed 2 | |||||||
| ID | L1 | Cos | SR | L1 | Cos | SR | L1 | Cos | SR |
| Ex.1 | 0.407 | 0.834 | 74.0 | 0.388 | 0.848 | 79.7 | 0.430 | 0.814 | 76.7 |
| Ex.2 | 0.367 | 0.861 | 77.2 | 0.353 | 0.874 | 82.7 | 0.387 | 0.845 | 80.3 |
| Ex.3 | 0.358 | 0.870 | 78.5 | 0.342 | 0.882 | 84.0 | 0.377 | 0.854 | 81.2 |
| Ex.4 | 0.389 | 0.853 | 80.6 | 0.372 | 0.867 | 86.0 | 0.409 | 0.835 | 82.9 |
| RGB-D Completion | Drawer Forecasting | Cup Forecasting | ||||
| ID | L1 | Cos | L1 | Cos | L1 | Cos |
| Ex.1 | 0.308 0.008 | 0.931 0.005 | 0.427 0.024 | 0.815 0.019 | 0.490 0.030 | 0.750 0.028 |
| Ex.2 | 0.263 0.006 | 0.947 0.004 | 0.387 0.020 | 0.847 0.016 | 0.456 0.026 | 0.785 0.023 |
| Ex.3 | 0.271 0.008 | 0.945 0.006 | 0.373 0.020 | 0.857 0.016 | 0.433 0.025 | 0.805 0.022 |
| Ex.4 | 0.314 0.008 | 0.928 0.006 | 0.400 0.021 | 0.842 0.017 | 0.456 0.027 | 0.786 0.025 |
| Meta-World | Drawer | Cup | ||||||
| ID | PL | MPC | Pull | Close | Full | Pick | Place | Full |
| Ex.1 | 70.9 3.2 | 83.8 2.4 | 82.2 3.8 | 85.6 1.9 | 65.6 5.1 | 80.0 3.3 | 81.1 1.9 | 61.1 5.1 |
| Ex.2 | 73.6 2.9 | 87.1 2.5 | 90.0 3.3 | 91.1 1.9 | 74.4 5.1 | 86.7 3.3 | 90.0 3.3 | 71.1 5.1 |
| Ex.3 | 74.9 2.6 | 88.1 2.7 | 91.1 3.8 | 93.3 3.3 | 77.8 5.1 | 88.9 3.8 | 91.1 1.9 | 74.4 5.1 |
| Ex.4 | 76.9 2.4 | 89.7 2.7 | 95.6 1.9 | 94.4 1.9 | 80.0 6.7 | 93.3 3.3 | 93.3 0.0 | 77.8 5.1 |
State Representation Learning
State representation learning develops compact, task-relevant descriptions of observations and histories for prediction, control, and decision-making.
See all 124 papers