JEPA world models can memorize backgrounds and forget actions: motion-guided contrastive regularization lifts planning success from 52% to 78% under randomized backgrounds
Synopsis
When a JEPA-style world model is regularized only through its embedding distribution, the encoder can treat the fixed background of each episode as its main content and suppress the state changes produced by robot-object interaction, a failure the authors call static collapse. The new method contrasts the predicted latent transition with transitions to other observations from the same episode, and uses optical-flow similarity as soft targets so candidates with similar motion receive similar supervision. Planning success rises from 52% to 78% over LeWM on OGBench-Cube with randomized backgrounds and from 14% to 54% on LIBERO-Object, while four other visual-control benchmarks stay comparable.
Interpretation
When visual backgrounds are complex, a JEPA-style world model can encode the background appearance that stays fixed within an episode into its latent state, compressing the robot and object state changes and degrading downstream planning. Earlier practice relied on distributional regularizers such as SIGReg, which push embeddings toward an isotropic Gaussian and constrain only the statistics, not which information is kept; LeWM reached 74% success on the original Cube and dropped to 52% on Cube-RB, which replaces only distant background pixels. Cube and Cube-RB share the robot, cube, actions and dynamics and differ only in the episode-specific background image, and both are trained with identical hyperparameters, so the success gap can be attributed to the background change; on t-SNE, Cube-RB embeddings cluster into tight episode-specific groups by background identity.
During training, the predicted next latent change is compared with latent changes to other observations from the same episode, so background identity can no longer separate action outcomes and the model is pushed to encode changes in the robot and manipulated objects. The contrastive target uses optical-flow similarity to build soft labels rather than treating every candidate as an equally dissimilar negative or measuring similarity by pixel differences; optical flow is computed offline with a pretrained SEA-RAFT model, and training reads only cached scalar distances. On Cube-RB and LIBERO-Object, adding plain InfoNCE raises success from 52% and 14% to 56% and 34%, and the motion-guided objective raises it further to 78% and 54%; removing SIGReg drops it back to 58% and 14%.
Perspective
The demonstrated scope is OGBench Cube and Cube-RB, LIBERO-Object, and the visually simpler Push-T, 2-Rooms and Reacher settings, all evaluated under a fixed sampling-based model-predictive control procedure; the method assumes static backgrounds and relies on offline data with sufficient action coverage. The authors propose improving transfer across actions and appearances, isolating task-relevant motion in dynamic scenes, and extending planning to longer horizons.
The authors state that, like LeWM, the method depends on sufficient action coverage and remains sensitive to shifts in action representations and visual appearance, including changes to LIBERO's rendering engine; background motion could distort the optical-flow similarities used for supervision. Hyperparameter sweeps show that adding candidates does not monotonically improve success, and a larger contrastive loss weight drops success to 38%.