EpicWorldModel: Exploration-driven Planning with Latent World Models

• Latent-space planning compresses sensory observations into structured representations for model-based planning, while JEPA predicts future or masked embeddings rather than reconstructing pixels.Prior JEPA world models faced collapse, while LeWorldModel stabilizes training by regularizing representations toward an isotropic Gaussian distribution [Maes et al., 2026].Source: §1 · p. 1
• Existing JEPA world models often use deterministic single-target latent regression, which can fail when history and actions do not uniquely determine the future.For example, a single predicted drawer content may mismatch the real contents and misdirect planning.Source: §1 · p. 2
• EpicWorldModel trains a stochastic JEPA to predict multiple plausible futures in a jointly learned representation space and plans under uncertainty.Its conditional flow-matching predictor uses past observations and actions as conditioning signals.Source: §1 · p. 2
• Predictive variance guides exploration: high variance signals uncertain future frames, while recently observed futures receive near-deterministic predictions.The paper relates flow predictive variance to an upper bound on predictive entropy and incorporates it into CEM planning.Source: §1 · p. 2
• EpicWorldModel achieves best or on-par performance among baselines, outperforming existing latent world models by up to 22% across partially observed tasks.The paper also reports state-of-the-art performance on visually restricted maze navigation tasks.Source: §1 · p. 2

• Latent world models target useful representations rather than pixel reconstruction, which can be computationally costly for real-time inference and control.Action-conditioned models such as DINO-WM and LeWorldModel apply learned representations to robot planning and control [Zhou et al., 2025, Maes et al., 2026].Source: §2 · p. 2
• Figure 1 contrasts a deterministic model that hallucinates a memorized fridge location with a stochastic model that maintains multiple plausible futures.Dark gray frames denote observation inputs and colored frames denote predicted futures; only four possibilities at t and one choice at t + 1 are shown.Source: Figure 1 · p. 3
• Flow matching reformulates denoising as an ODE, and EpicWorldModel uses a flow ODE to predict future representations within the JEPA paradigm.The related-work discussion traces flow matching from continuous normalizing flows and conditional flow matching [Chen et al., 2018, Lipman et al., 2022].Source: §2 · pp. 2–3
• Prior uncertainty-aware planning has used prediction error, probabilistic ensembles, and ensemble disagreement as intrinsic signals for exploration [Pathak et al., 2017, Chua et al., 2018, Sekar et al., 2020].EpicWorldModel draws on this exploration literature while using a stochastic JEPA for both goal-reaching and curiosity signals.Source: §2 · p. 3

3 EpicWorldModel

• The framework predicts a conditional distribution over future latent states, supporting plans across multiple plausible futures when critical scene content is occluded or absent.It jointly learns the stochastic predictor and latent representation space.Source: §3 · p. 3

3.1 Problem Formulation

• In a POMDP, hidden states generate observations and actions drive transitions, so the agent plans from a finite observation-action history.The formulation is goal-conditioned and includes a reward and discount factor.Source: §3.1 · pp. 3–4

3.2 Learning the Stochastic Latent Predictive World Model

• Conditional flow matching transports Gaussian source samples toward future latent observations conditioned on history and action chunks, while SIGReg regularizes embeddings toward an isotropic Gaussian.The complete objective jointly supervises stochastic prediction and a well-defined latent representation.Source: §3.2 · pp. 4–5
• Validation losses show smooth convergence, which the paper attributes to the bounded gradient noise established by Proposition 1.Figure 2 shows smoothed and per-step loss curves for visual PointMaze.Source: Figure 2 · p. 5, §3.2 · p. 6

3.3 Uncertainty-Aware Latent Space Planning

• The planner uses disagreement across sampled future predictions as an uncertainty signal and incorporates it into CEM to balance exploration and goal-reaching.The paper motivates this signal as a proxy that approximately maximizes an upper bound on expected information gain.Source: §3.3 · p. 6
• EpicWorldModel achieves the best performance in both default- and long-horizon settings across benchmark environments.Figure 3 compares the FM and iMF flow-matching objectives.Source: Figure 3 · p. 7

4.1 Baselines

• Comparisons include goal-conditioned RL baselines GCIQL and GCIVL, plus world-model planners PLDM, DINO-WM, and LeWorldModel.The experiments cover OGBench navigation tasks, Car Racing, LabMaze, and RoboCasa NavigateKitchen.Source: §4.1 · p. 7, §4 · p. 7

4.2 Implementation Details

• History tokens combine action and temporal embeddings, future tokens receive temporal conditioning, and causal masking limits future tokens to history and earlier future tokens.Inference samples 10 noise samples per action chunk and uses either 2-NFE average velocity or 10-NFE FM; β and NFE are selected on held-out scenes.Source: §4.2 · pp. 7–8
• Training uses heuristic-expert trajectories for OGBench and LabMaze, and 500 human demonstrations for RoboCasa NavigateKitchen.The reported training time ranges from under 5 GPU hours for Visual PointMaze to up to 20 GPU hours for Visual AntMaze.Source: §4.2 · p. 8

4.3 Planning Performance of EpicWorldModel

• On Visual PointMaze Giant, FM reaches 86% ± 2.6 success versus LeWM’s 64.5 ± 2.9; performance is on-par on AntMaze and competitive on Visual Scene, Car Racing, and LabMaze.The authors also report substantially better results than goal-conditioned policy baselines on PointMaze.Source: §4.3 · pp. 8–9, Figure 5 · p. 8

4.4 Planning Performance Under Controlled Uncertainty Levels in RoboCasa

• At 1.26 m, exploration raises iMF success from 46.7% to 80.0%; FM reaches 93.3% at the largest distance versus 20.0% for LeWM.Under matched planning compute, the stochastic predictor without exploration achieves 71.7% pooled success versus 45.0% for deterministic LeWM.Source: §4.4 · p. 9, Table 1 · p. 10
• With initial goal visibility below 0.5, exploration improves success from 1/6 to 5/6; at visibility ≥0.95, it reduces success from 21/22 to 19/22.The experiments provide preliminary evidence that uncertainty-driven exploration yields information about hidden goal location.Source: §4.4 · p. 10, Table 2 · p. 11

4.5 Ablation Study on the choice of β and the Number of Flow Evaluations (NFE)

• Across most tested environments, β = 0.4 and NFE = 2 offer a favorable exploration–compute trade-off, while Visual Scene favors NFE = 1.The authors leave automatic adaptation of β to future work.Source: Figure 6 · p. 8, §4.5 · p. 11

4.6 Analysis on Exploration Behavior in Latent Space Planning

• In one rollout, the top-k covariance ellipse shrinks from 60.99 to 11.86 as the agent approaches the goal, and the mean prediction points toward it.This qualitative example is consistent with predictive uncertainty decreasing near the goal.Source: Figure 7 · p. 9, §4.6 · p. 11

5 Conclusion

• EpicWorldModel is an uncertainty-aware planning framework using stochastic JEPAs based on flow matching.The conclusion reports theoretical motivation and experimental validation, with high planning success rates, especially in partially observable environments.Source: §5 · p. 11
• Goal-image conditioning can be restrictive because real-world tasks may be specified through language and goal images may provide privileged deployment information.The conclusion identifies alignment with language-based or abstract goal specifications as future work.Source: §5 · p. 11
• The authors report that EpicWorldModel is a stable JEPA and demonstrates promise for exploration-driven planning with latent world models.Source: §5 · p. 11

A Proofs

• Under an isotropic Gaussian latent space, the conditional velocity target has bounded, data-independent excess variance that does not grow with representation expressiveness or conditioning complexity.The proof derives the result using the law of total variance and the conditional expectation of the velocity target.Source: A Proofs · p. 15
• The Gaussian latent assumption supports the variance analysis: the proof treats the source and target latents as jointly Gaussian, yielding homoscedastic conditional variances.The conditional-variance derivation is stated for each flow time.Source: A Proofs · p. 15
• For a well-trained predictor, U is an unbiased estimator of mean conditional variance for an action-chunk embedding.The proof uses the sample variance across predicted futures as an unbiased estimate of marginal variance, averaged over dimensions.Source: A Proofs · pp. 15–16
• The maximum-entropy argument upper-bounds predictive entropy using U, whose scale ranges from concentrated predictions near 0 to prior-diffuse predictions near 1.The bound is tight only when the predictive covariance is isotropic; for anisotropic covariances it may be loose, so U is used as a practical exploration bonus rather than an entropy estimate.Source: A Proofs · p. 15

B Additional Implementation Details

• The visual encoder maps 224 × 224 observations to 192-dimensional latents, and the same encoder embeds goals for shared-space distance calculations.The encoder is a ViT-Tiny with patch size 14, trained from scratch; a two-layer MLP with BatchNorm projects its CLS token.Source: B Additional Implementation Details · p. 16
• With frame skip F=5, action vectors (10-dimensional in PointMaze, where da=2) are embedded into one D-dimensional action token per predicted future step; output returns only K future-step predictions.The transformer processes history and future tokens, then slices the output to out[:, K:] for the predicted velocity field.Source: B Additional Implementation Details · p. 16
• LeWM, iMF, and FM share a CEM planner over 25-step action sequences, while EpicWorldModel scores sampled rollouts using mean goal error minus βU.The planner runs 30 iterations with 300 candidates and executes the best mean open-loop before replanning.Source: B Additional Implementation Details · p. 16
• RoboCasa variants share data, encoder, projector, and hyperparameters, differing only in predictor family; evaluations use final epoch-100 checkpoints.Training uses 500 NavigateKitchen demonstrations and a frozen DINOv3 encoder; the compared predictors are deterministic AR, one-step iMF, and 10-step FM.Source: B Additional Implementation Details · pp. 16–17

C Additional Results

• The additional experiments are presented as complements to the main results.Source: C Additional Results · p. 17
• Additional experiments include disabling exploration bonus U, a compute-matched deterministic-baseline comparison, a sampled-futures sweep, and a population-level latent-space check.Source: C Additional Results · p. 17
• Unless otherwise noted, the additional experiments use the same CEM/MPC planning protocol as the main paper.Source: C Additional Results · p. 17

C.1 Effect of the Exploration Term (β = 0)

• Disabling exploration lowers success by 7 points for iMF at NFE = 4 and by 6 points for FM at NFE = 10, both with β = 0.4.The comparison uses a β = 0 condition.Source: §C.1 · p. 17, Table 3 · p. 17
• The stochastic predictor alone explains part, but not all, of the gain over deterministic LeWM.The remaining gain is attributed to the exploration term in the planning cost.Source: §C.1 · p. 17
• Table 3 reports success rates from a single training and evaluation run per cell.It compares disabling exploration across NFE settings on Visual PointMaze Giant.Source: Table 3 · p. 17

C.2 Compute-Matched LeWorldModel

• LeWM’s success on PointMaze rises from 54.0% to 74.0% with more action candidates, but remains about 10 points below EpicWorldModel iMF.The remaining gap persists at lower and matched compute.Source: §C.2 · p. 17, Table 4 · p. 18
• Scaling LeWM’s action-candidate budget recovers only part of the gap, even when compute is matched to stochastic rollouts.The comparison approximates N × NFE = 20 evaluations with up to 6,000 LeWM candidates.Source: §C.2 · p. 17, Table 4 · p. 18
• LeWM planning compute does not explain AntMaze success, which is described as bounded by task and action-space complexity.The passage refers to the discussion of AntMaze success rates in Section 4.3.Source: §C.2 · p. 17

C.3 Number of Sampled Futures

• Across all four Visual PointMaze settings, changing the number of sampled futures N has only a small effect on performance.The sweep tests N ∈ {2, 5, 10} at β ∈ {0, 0.4} for FM and iMF.Source: §C.3 · p. 18, Table 5 · p. 18
• The sweep varies N separately from planning horizon K, with N = 5 as the default used in the main results.FM uses NFE = 10 and iMF uses NFE = 2 in this Visual PointMaze ablation.Source: §C.3 · p. 18, Table 5 · p. 18
• Automatic selection of N beyond a single default across tasks is left for future work.Source: §C.3 · p. 18

C.4 Latent-Space Analysis over All Evaluation Scenarios

• The top-k latent predictions contract as the agent approaches the goal, with 1σ ellipse area shrinking from 60.99 to 11.86.Source: §C.4 · p. 18
• Across 50 scenarios, median goal progress reaches 45.5%; distance to the goal latent falls from 17.4 ± 5.1 at t = 0 to 10.4 ± 5.8 at t = 25.Median progress is 58.6% among the 80% of episodes that succeed.Source: §C.4 · p. 18
• Among 1,500 sampled future predictions, the closest sample has median distance 3.5 to the ground-truth future latent and median cosine alignment 0.88 with its direction.The authors interpret these results as futures tracking the ground-truth trajectory rather than producing arbitrary spread.Source: §C.4 · p. 18