Search papers, labs, and topics across Lattice.
This paper introduces an end-to-end JEPA world model that enhances action-conditioned latent prediction with inverse dynamics and state alignment to improve robotic planning towards visually specified goals. The integration of inverse dynamics prevents latent collapse and ensures that learned representations are informative about the actions taken, while state alignment grounds these representations in their physical context. The proposed model achieves superior success rates across multiple benchmark tasks, demonstrating the effectiveness of state alignment in enhancing planning performance compared to existing methods like LeWorldModel.
State alignment significantly boosts planning success in robotic tasks, achieving 100% success on TwoRoom and 98% on PushT, outperforming traditional models.
Action-conditioned JEPA world models enable planning toward visually specified goals without reconstructing future pixels, yet latent prediction alone does not explicitly encourage the learned representations to retain information relevant to robotic control. We introduce an end-to-end JEPA world model that augments latent prediction with inverse dynamics (IDM) and state alignment (SA). While inverse dynamics discourages latent collapse and makes latent transitions informative of the actions that produced them, state alignment grounds consecutive representations in their associated physical configuration and motion. Across four benchmark tasks, our model attains the highest success rates on TwoRoom (100%), PushT (98%), and OGBench-Cube (87%), while performing comparably to LeWorldModel on Reacher. Our ablation further shows that adding state alignment consistently improves planning success over IDM alone across all four tasks. Although LeWorldModel, our primary baseline, attains higher average straightening on OGBench-Cube, transition-subspace analysis shows that its transition energy is concentrated in a substantially lower-dimensional subspace. Our state-aligned model exhibits a higher effective transition dimension than LeWorldModel and improves planning over IDM alone, supporting state alignment as an effective complement to inverse dynamics for robotic planning.
