EpicWorldModel: Exploration-driven Planning with Latent World Models

AI summary

Overview

Research area: Machine learning, specifically latent world models, model-based reinforcement learning, and robot planning under partial observability.

Technical level: Advanced. The paper assumes familiarity with joint-embedding predictive architectures (JEPAs), flow matching, POMDPs, and information-theoretic quantities such as entropy and expected information gain. The plain-English methodology section below makes the core ideas accessible without that background.

Scope: The paper proposes EpicWorldModel, a stochastic JEPA-based latent world model trained with a flow-matching objective, and shows that the variance of its predictions can be used as an exploration bonus inside CEM-based planning to reach occluded goals.

What This Paper Is About

Existing JEPA world models predict a single future latent state, which works when the future is determined by what the agent has already seen. Under partial observability, however, many futures are consistent with the same observation history and action. This paper asks how to learn a world model that predicts multiple plausible futures when scene content relevant to the goal is occluded, and how to use the model's own predictive uncertainty to guide a planner toward regions where hidden goals are likely to be.

Key Contributions

1. A stochastic JEPA world model. EpicWorldModel replaces the deterministic single-target regression of prior latent world models with a conditional flow-matching predictor that samples multiple plausible future latent states given a latent history and an action chunk. The encoder and predictor are learned jointly, and an average-velocity formulation collapses inference to a single function evaluation (1-NFE).

2. Flow predictive variance as an exploration signal. The authors argue, and formalize, that when the model is uncertain the flow-matching predictor shows high variance across generated futures, and that this variance relates to an upper bound on predictive entropy and to expected information gain. They use ensemble disagreement across sampled futures as a scalar uncertainty measure U.

3. Uncertainty-guided planning. The uncertainty signal is folded into Cross-Entropy Method (CEM) planning as a bonus term weighted by β, so the planner balances reaching the goal against exploring uncertain regions where occluded goals are most likely located.

4. Empirical evaluation across partially observable benchmarks. The method is evaluated on Visual PointMaze Giant, Visual AntMaze Giant, Visual Scene, Car Racing, LabMaze, and RoboCasa NavigateKitchen, reporting best or on-par performance and up to 22% improvement in success rate over LeWorldModel.

Main Findings

• Visual PointMaze Giant: EpicWorldModel achieves the best performance across all methods, with approximately 22% higher success rate than LeWorldModel. The iMF variant reaches 84% ± 3.6 and the FM variant 86% ± 2.6, compared with LeWorldModel at 64.5 ± 2.9 and DINO-WM at 67.8 ± 1.8.

• Visual AntMaze Giant: The task is harder because of higher-dimensional action spaces. FM reaches 30% success and iMF 28%, which is on par with LeWorldModel, DINO-WM, and PLDM.

• Car Racing, LabMaze, Visual Scene: On Car Racing, iMF and FM achieve 30% and 28% success respectively, roughly double LeWorldModel (16.0%) and above DINO-WM (24.0%); PLDM fails out-of-regime. On LabMaze, iMF again performs best. On Visual Scene the methods remain competitive but the gap is smaller, which the authors attribute to flow-based latent prediction being most effective in partially observable tasks with relatively low-dimensional action spaces.

• Long-horizon planning: Doubling the target planning horizon consistently reduces performance across environments.

• Gains are not explained by inference compute: Giving LeWorldModel 6000 CEM candidates (20× the default) raises PointMaze success only to 74.0%, still approximately 10 points below iMF. On AntMaze, extra samples do not help (24.0–26.5%).

• RoboCasa controlled difficulty: Under matched planning compute, the stochastic predictor without exploration outperforms deterministic LeWorldModel, 71.7% versus 45.0% success pooled across start distances. Increasing LeWorldModel to 3× more candidates improves short-range performance but provides little benefit at larger distances. At 1.26 m, enabling exploration increases iMF success from 46.7% to 80.0%, while leaving easier short-range settings unchanged. FM performs best at the largest distance, reaching 93.3% versus 20.0% for LeWorldModel.

• RoboCasa controlled occlusion: The uncertainty bonus helps most when the goal is initially occluded (visibility < 0.5), improving success from 1/6 to 5/6. The benefit decreases as visibility increases and can reduce goal-reaching when the goal is fully observable (≥ 0.95, going from 21/22 to 19/22). Pooled across visibility strata, LeWorldModel achieves 13/39 (33.3%), iMF with β = 0 achieves 29/39 (74.4%), and iMF with β = 0.1 achieves 31/39 (79.5%).

• Ablation on β and NFE: Visual PointMaze Giant shows clear sensitivity to hyperparameters, with optimal performance at β = 0.4 and NFE = 2. AntMaze Giant is more stable across settings with modest gains from moderate exploration. Visual Scene performs better with NFE = 1. (The text of this ablation section is truncated in the provided content; results are shown in Figure 6.)

• Theory on the latent geometry: With the isotropic Gaussian latent regularization, the excess variance of the conditional velocity target is 1 / ((1 − τ)² + τ²) · I_d, depending only on latent dimensionality and flow step τ, so gradient variance is bounded and data-independent. The uncertainty measure U is an unbiased estimator of the mean conditional variance, tr(Σ_{c_t}) / d, and the predictive entropy is upper-bounded by (d/2) log(2πe · U). U is bounded within [0, 1]. The bound is tight only for Σ_{c_t} ∝ I_d, so the authors use U as a practical exploration bonus rather than an entropy estimate.

• Training behavior: Loss curves on the visual PointMaze environment show minimal variance and smooth convergence of the encoder and stochastic predictor.

Methodology in Plain English

The system has two learned parts. An encoder turns each observation into a compact latent vector, and a predictor models how those latent vectors evolve when actions are taken. Unlike earlier world models, which regress one future vector, this predictor is generative: it starts from a random Gaussian noise vector and transports it to a future latent state.

The transport is done with conditional flow matching. Training constructs a straight-line interpolation between a noise sample and the true future latent block, and trains the network to predict the velocity along that line. To keep planning fast, the authors use an average-velocity formulation (improved MeanFlow) so the network can predict the total displacement over a time interval in one step, allowing inference with a single function evaluation by choosing τ = 0 and r = 1. The latent space is simultaneously regularized toward an isotropic Gaussian using the Sketched Isotropic Gaussian Regularizer (SIGReg), which projects embeddings onto random unit-norm directions and applies the Epps-Pulley normality test to each projection. The full objective is the flow-matching loss plus a weighted SIGReg term. Matching the noise source to the latent marginal makes transport paths shorter and better conditioned.

At planning time, the model draws 10 noise samples per action chunk and runs the predictor on each, producing an ensemble of predicted futures. The disagreement between ensemble members, averaged over latent dimensions, is the uncertainty score U. Planning uses Model Predictive Control with the Cross-Entropy Method: the action chunk is chosen to minimize the sum over samples of the squared latent distance to the goal embedding, minus β times U. A positive β therefore encourages the planner to prefer action chunks whose futures are uncertain, which in the RoboCasa experiments corresponds to turning and moving toward parts of the kitchen that have not yet been observed. The predictor is implemented as a transformer over a sequence of H history tokens and K future tokens with causal attention masking, so each future token can only attend to history tokens and earlier future tokens.

Why This Matters

For research, the paper identifies a concrete failure mode of deterministic JEPA world models under partial observability — memorizing and hallucinating goal locations seen during training — and offers a stochastic alternative with an information-theoretic justification. It also shows that a purely latent model can succeed on visually restricted maze navigation tasks that were previously challenging for latent world models, and it connects flow-matching variance to entropy and expected information gain, which may be reusable in other generative-model-based planners.

Real-world applications suggested by the evaluation environments:

• Home and kitchen service robots that must find appliances such as a fridge that is outside the field of view at the start of an instruction (RoboCasa NavigateKitchen).
• Warehouse and logistics robots navigating cluttered aisles where destinations are blocked by shelving or equipment.
• Autonomous driving in traffic where other agents and obstacles are intermittently occluded (evaluated on Car Racing).
• Search-and-rescue or exploration robots that must actively look into unobserved regions before a target becomes visible (LabMaze-style navigation).

For industry, the work is relevant to companies building embodied agents and mobile manipulators, since the authors report evaluating 50 test scenarios in parallel on a single A6000 GPU with 48 GB of GPU-RAM and training times as low as under 5 GPU hours for Visual PointMaze and up to 20 GPU hours for Visual AntMaze with the iFM objective. One co-author is affiliated with Torc Robotics, indicating direct relevance to autonomous systems. The key industrial argument is that exploration under uncertainty is a capability that cannot be substituted with more planning compute: extra CEM candidates did not close the gap for the deterministic baseline.

Future Directions

• **Long

Authors’ 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.