# SWIM: Learning State Weights from Behavior to Train JEPA World Models

- Source: arXiv cs.RO
- Published: 2026-10-28T04:00:00.000Z
- RoboHot: https://robohot.news/en/items/a5uyzc5yn0evcu1tj4blsjkfy
- Original: https://arxiv.org/abs/2610.36662

## Summary

SWIM proposes using an inverse dynamics model to learn state feature weights from offline trajectories for state-supervised training of JEPA world models. The method predicts recorded actions from current states and observed future displacements, quantifies each state feature's contribution to action prediction, and uses this to build a weighted state metric that aligns the latent space with state distances. The inverse dynamics model is used only during training; planning still relies solely on images. On five simulated robotic manipulation tasks, SWIM's planning success rate consistently outperforms uniform state weighting.