Motion-guided Contrastive Regularization For Latent World Models

Yiwen Chen; Qiqi Hou; Jisoo Jeong; Yinhao Zhu; Huaizu Jiang; Amir Said

cs.CVNov 4, 2026

Motion-guided Contrastive Regularization For Latent World Models

Authors: Yiwen Chen, Qiqi Hou, Jisoo Jeong, Yinhao Zhu, Huaizu Jiang, Amir Said

Organizations: Northeastern University · Qualcomm AI Research

Abstract

JEPA-style world models learn action-conditioned latent dynamics without reconstructing observations, but distributional regularization alone does not ensure that their representations preserve information needed for control. We identify static collapse, a failure mode in which the encoder preserves episode-specific static content while suppressing manipulation-relevant state changes. Because static content is temporally predictable yet diverse across episodes, this failure can persist despite regularizers such as SIGReg. A controlled study on OGBench-Cube shows that introducing randomized backgrounds substantially degrades planning performance while leaving the underlying manipulation task unchanged. To address this problem, we propose motion-guided contrastive regularization. Our approach compares predicted latent transitions with transitions to candidate observations from the same episode, reducing the usefulness of static background identity for distinguishing action outcomes. Optical-flow similarities provide soft contrastive targets, allowing candidates with similar motion to receive similar supervision. Combined with latent prediction and SIGReg, our objective encourages representations to preserve within-episode state distinctions and similarities between related transitions. Our method improves planning success over LeWorldModel (LeWM) from 52% to 78% on OGBench-Cube with randomized backgrounds and from 14% to 54% on LIBERO-Object, while maintaining comparable performance on four other visual-control benchmarks.

Figures & tables

Figure 1: Static collapse and our solution. (A) We train JEPA world models end-to-end on OGBench-Cube (Cube) and a variant with randomized backgrounds (Cube-RB). In Cube-RB, one background image is sampled per episode and held fixed throughout it. Only distant pixels (depth $\geq 1.0$ m) are replaced, so the robot, cube, and dynamics are unchanged and only task-irrelevant background variation is added. (B) t-SNE visualizations of the encoder embeddings, where each point is one input frame and each color one episode. On Cube-RB, LeWorldModel ( Maes et al., 2026 , LeWM;) maps distinct states within each episode to tightly clustered latents, whereas our motion-guided contrastive regularization preserves within-episode state variation. (C) Planning success of LeWM, Planning with Latent Dynamics Models ( Sobal et al., 2026 , PLDM;) , and our method. LeWM and our method perform similarly on Cube, but LeWM’s success drops on Cube-RB and is low on LIBERO-Object. Our method achieves the highest success of the three in all settings.

Figure 2: Motion-guided contrastive learning. The world model encodes the current observation $\mathbf{x}_{t}$ and predicts $\hat{\mathbf{z}}_{t+1}$ . The true next observation $\mathbf{x}_{\mathrm{pos}}=\mathbf{x}_{t+1}$ provides the positive sample, while other frames from the same episode provide negative candidates $\mathbf{x}_{m}$ and $\mathbf{x}_{n}$ . All observations share the same encoder. Image-space: Optical flow is computed from $\mathbf{x}_{t}$ to each candidate. Comparing these fields with the positive flow assigns higher target similarity to candidates with similar motion ( $\mathbf{x}_{m}$ ) than to those with different motion ( $\mathbf{x}_{n}$ ). Latent-space: Angles illustrate cosine similarities between the predicted transition $\hat{\mathbf{z}}_{t+1}-\mathbf{z}_{t}$ and candidate transitions $\mathbf{z}_{j}-\mathbf{z}_{t}$ . Our motion-guided contrastive loss matches the resulting latent similarity distribution to the motion-derived target distribution through cross-entropy. The prediction loss aligns $\hat{\mathbf{z}}_{t+1}$ with $\mathbf{z}_{\mathrm{pos}}$ . The grippers in the sample frames are highlighted for illustration purposes only.

Method	Push-T	2-Rooms	Reacher	Cube	Cube-RB	LIBERO- Object	Peak GPU mem. (MB)	Planning time (ms)
DINO-WM † ( Zhou et al., 2024 )	74	100	78	78	68	68	20,714	11,223
PLDM ( Sobal et al., 2026 )	78	97	64	65	62	38	338	28
SMWM ( Ivashkov et al., 2026 )	83	99	66	84	78	40	374	101
LeWM ( Maes et al., 2026 )	90	86	80	74	52	14	372	112
Ours	90	88	78	76	78	54	372	112

Table 1: Planning success rate (%) on six settings. † uses a frozen pretrained visual encoder. Peak GPU memory and planning time are measured on Cube dataset with an A5000 GPU and averaged over its 50 evaluation instances; planning time is the wall-clock time to plan one full trajectory (five replanning rounds). CEM uses 128 candidates and 10 iterations per round, with a planning horizon of up to 5 control steps.

Panel of Figure 3 — Figure 3: Latent-space decoding on Cube, Cube-RB, and LIBERO-Object. Compared to LeWM, reconstructions from our frozen latents more clearly preserve the robot and manipulated objects on the challenging Cube-RB and LIBERO-Object datasets. Check our video demo for more comparisons.

Panel of Figure 4 — Figure 4: Encoder attention over image patches (brighter = more attention). On Cube-RB and LIBERO-Object, LeWM attends to background and static objects; ours focuses on the gripper and manipulated object.

Dataset	Metric	LeWM	Ours
Cube	Robot pos. error (cm) $\downarrow$	2.74	0.61
Cube	Cube pos. error (cm) $\downarrow$	1.27	0.62
Cube-RB	Robot pos. error (cm) $\downarrow$	13.01	3.56
Cube-RB	Cube pos. error (cm) $\downarrow$	10.49	5.10
LIBERO- Object	Robot pos. error (cm) $\downarrow$	10.67	6.00
LIBERO- Object	Holding object (AUROC) $\uparrow$	0.78	0.88

Table 2: MLP probing of physical state from frozen latents. Position errors are Euclidean distances in centimeters. Holding object (AUROC) is a binary gripper-state classification on LIBERO-Object.

SIGReg	Contrastive objective	Cube-RB	LIBERO-Object
$\checkmark$	None	52	14
$\checkmark$	InfoNCE	56	34
–	Motion-guided	58	14
$\checkmark$	Motion-guided	78	54

Table 3: Ablation of training objectives on Cube-RB and LIBERO-Object. Entries are planning success (%). All variants retain the prediction loss. InfoNCE uses the same within-episode candidates (

N=6

) and latent transition scores (

\tau_{z}=0.1

) as our method.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

(a) Flow temp.		(b) Latent temp.		(c) Candidates		(d) SIGReg		(e) Contrastive
$\tau_{f}$	Succ.	$\tau_{z}$	Succ.	$N$	Succ.	$\lambda$	Succ.	$\beta$	Succ.
2.5	40	0.02	52	2	52	0	14	0.05	50
5	36	0.05	54	4	48	0.03	40	0.1	54
10	46	0.1	54	6	54	0.09	54	0.3	54
20	54	0.2	48	8	50	0.18	54	0.6	52
40	40	0.5	50	10	54	0.3	50	1.0	38

Appendix

Table 4: Hyperparameter sensitivity on LIBERO-Object. “Succ.” denotes planning success (%). Bold marks the highest result within each sweep; underlining marks the reference parameter value.

World Models for Robotics

World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.

See all 405 papers

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models
Aug 18, 2026Jack Boylan, Chris HokampJoint-Embedding Predictive Architecture Representation Collapse
Quantexa
Learning Counterfactual World Models for Embodied Reasoning under Partial Observability
Sep 5, 2026Todd Y. Zhou, Daniel ZhangEmbodied World Models World Model-Based Planning
Harvard University
Anisotropic Representations Improve Planning in JEPA World Models
Sep 29, 2026Mingu Kang, Yoori Oh, Sookyung Kim +1Representation Learning World Model-Based Planning
Seoul National University · Ewha Womans University

Aug 18, 2026J/K move · Enter open · S save

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

Jack Boylan, Chris Hokamp

Quantexa

Joint-Embedding Predictive Architecture Representation Collapse Action-Conditioned World Models Inverse Dynamics

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future embeddings, but the objective admits a trivial solution of a constant encoder, so every practical system adds an anti-collapse mechanism (LeCun, 2022; Assran et al., 2023; Bardes et al., 2022; 2024). LeWorldModel (LeWM) prevents collapse with SIGReg, a regularizer that forces the latent distribution to match an isotropic Gaussian: the representation is stabilized by prescribing what it must look like, independently of the environment it models. We argue that the anti-collapse pressure can instead come from the transition data itself. Action-Contrastive Masked Transition Modeling (AC-MTM) keeps LeWM's forward latent-prediction objective and adds a training-only inverse-dynamics head trained with Action-NCE: each latent transition must identify the action that produced it among the other actions in the batch, a discrimination task that a collapsed encoder provably fails. The inverse branch is discarded after training, leaving test-time encoding, forward prediction, planning, and compute identical to LeWM. On four standard pixel-control tasks under a matched planning protocol, AC-MTM trains stably from scratch and matches SIGReg on average. On the harder multi-object OGBench Visual Scene task, results are consistent with the prescribed geometry becoming a bottleneck: AC-MTM reaches 80.0

\pm

2.0% success versus 58.0

\pm

2.0% for SIGReg, improving by 20-24 points in each training seed. A single 50-episode random-policy run gives a 52% baseline estimate. Contrastive inverse dynamics thus provides a distribution-free anti-collapse signal that requires no target network, stop-gradient, pretrained encoder, or reconstruction objective, and we characterize the action-space and observability assumptions under which it holds. We make our code available at https://github.com/jackboyla/action-contrastive-jepa

Abstract

World Models for Robotics

Explore similar work

No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models

Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

Anisotropic Representations Improve Planning in JEPA World Models