EpicWorldModel: Exploration-driven Planning with Latent World Models

2026-10-05

AI Summary

This study addresses the limitation of deterministic latent world models in capturing multiple future possibilities under partial observability by proposing the Stochastic JEPA framework. This approach integrates Joint Embedding Predictive Architectures with flow matching to predict diverse future states, optimized via cross-entropy loss. Furthermore, it introduces a novel exploration mechanism based on flow prediction variance, leveraging uncertainty signals to guide agents in balancing goal achievement with the exploration of occluded regions. Experimental results demonstrate that the proposed framework achieves state-of-the-art or comparable performance across multiple tasks, improving success rates by up to 22% over baselines.

Abstract

Latent world models based on Joint-Embedding Predictive Architecture (JEPA) are deterministic by design. While successful in fully observable scenarios, this paradigm breaks down when past observations and actions lead to multiple plausible future possibilities, e.g., due to occlusion. We introduce EpicWorldModel, a framework to train stochastic JEPAs for environments and tasks with inherent uncertainty under partially observability. We jointly train the EpicWorldModel predictor with its latent representation space to directly predict multiple potential future states using a flow-matching objective, when the goal-relevant scene content is absent from the conditioning history. We show that flow predictive variance, motivated by its relation to an upper bound on predictive entropy, serves as a useful exploration guidance for planning. By incorporating this uncertainty signal into Cross-Entropy Method (CEM)-based planning, our approach balances goal-reaching with exploration of uncertain regions where occluded goals are most likely to be located. We demonstrate the effectiveness of EpicWorldModel through a series of latent planning experiments with the best or on-par performance across tasks, showing up to 22% empirical improvement in success rate over LeWorldModel.

Problem

• Latent World Models
• Partial Observability
• Uncertainty
• JEPA
• Planning

Innovation

• Latent World Models
• Stochastic JEPA
• Flow Matching
• Exploration-driven Planning
• Partial Observability