PRISM: PRior-guided Imagination Sampling in world Models

Abstract: A learned world model provides a powerful physical intuition for evaluating future states. But its effectiveness in continuous control also depends critically on how candidate actions are generated for model-based planning. Rather than solely asking how accurately a model can simulate the future, we ask: which candidate actions are worth evaluating in the first place? Existing planners typically search arbitrarily or use expert demonstrations only to initialize a sampling mean, discarding the expert's state-conditioned confidence. Properly guiding this search requires a robust action prior, yet current approaches often rely on independent visual encoders or large-scale VLMs to obtain one. We argue that this architectural bloat is unnecessary: the exact same data - and the learned representations of the world model itself - inherently encode the agent's action intuition. We introduce PRISM, a task-agnostic framework that extracts both from a single dataset while maintaining strict architectural simplicity. Building on a standard JEPA-style latent world model, PRISM attaches a lightweight MLP directly to its frozen encoder to predict a state-conditioned Gaussian prior. At plan time, PRISM fuses this prior into the planner's sampling distribution via a precision-weighted Product-of-Gaussians update. This parameter-free, closed-form integration steers the sampling process, making the prior confident where it is and ceding control where it is not. PRISM improves success rates by 35 percentage points over vanilla world-model-based MPC on Cube and 32 percentage points on PushT, without introducing significant inference overhead.

• The paper introduces a novel framework that integrates state-conditioned action priors into model predictive control via precision-weighted fusion.
• It demonstrates significant empirical improvements in visual manipulation tasks, increasing success rates by over 30 percentage points while preserving sample efficiency.
• The method transitions seamlessly to real-world robotics, ensuring robust performance with graceful degradation when prior confidence decreases.

PRISM: Prior-Guided Imagination Sampling in World Models

Motivation and Problem Definition

Planning in continuous control tasks with world models depends not only on the accuracy of the dynamics prediction but critically on the action sampling that initiates each planning episode. Existing world model planners typically rely on uninformed Gaussian distributions or use external expert demonstrations to warm-start sampling, discarding state-conditioned uncertainty. Current solutions either employ architectural complexity (such as vision-LLMs or independent encoders) to obtain action priors, incurring computational and perceptual redundancy, or lack the robustness mechanisms necessary to avoid planner derailment when the prior is unreliable. The central thesis is that the agent’s action intuition is encoded inherently in the latent representations learned by the world model itself, and can be extracted without auxiliary models.

Methodology and Architecture

PRISM introduces a task-agnostic, lightweight framework for incorporating an action prior into model predictive control (MPC) planners paired with Joint-Embedding Predictive Architecture (JEPA) world models. The approach attaches a 3-layer MLP prior head directly to a frozen JEPA encoder. This prior head, trained offline, predicts a state-conditioned unimodal Gaussian over the next horizon-length action block from the same demonstration dataset used by the world model predictor.

At plan time, PRISM integrates the prior via a closed-form precision-weighted Product-of-Gaussians (PoG) fusion with the planner’s default initialization. This fusion updates the sampling distribution’s mean and variance, ensuring that the prior influences the planner only proportionally to its predicted confidence for each action coordinate. This mechanism guarantees graceful degradation: when the prior is uncertain (high variance), the planner reverts to vanilla uninformed sampling automatically, with no learned gates or heuristics. PRISM exploits MPPI’s design, which maintains fixed covariance throughout optimization, allowing the prior's confidence to persist structurally during the planning loop.

The prior head is trained with a stabilized 𝛽-NLL objective, which interpolates between Gaussian NLL and MSE to prevent variance collapse and improve training stability, and uses hindsight goal sampling so that goal labels are self-consistent with demonstration outcomes.

Experimental Results

PRISM is evaluated on two visual manipulation tasks—PushT (planar T-block alignment) and Cube (single-arm cube manipulation)—both requiring the planner to match a target goal observation. The method consistently outperforms both planning and policy-only baselines, including behavior cloning and Diffusion Policy, across all sample budgets:

• PRISM-MPPI increases success rates by up to 35 percentage points on Cube and 32 points on PushT compared to vanilla world model MPC.
• Even with a minimal sample budget (e.g., $K=32$), PRISM-MPPI outperforms vanilla MPPI at high budgets ($K=128$), demonstrating superior sample efficiency.
• The integration of the prior’s per-state variance is pivotal: mean-only warm-starting (without variance) fails to scale reliably with increased sample budgets, underscoring the necessity of including precision.
• PRISM maintains computational parsimony: the additional inference overhead is microsecond-level, indistinguishable from baseline variance and compatible with real-time planning requirements.

Ablation studies confirm that PoG fusion is robust across multiple encoder architectures and prior parameterizations. PRISM retains its relative gains even when absolute task performance drops using a non-tailored DINOv2 encoder. The method is resistant to prior mis-tuning, smoothly degrading to the prior-free baseline as the prior variance increases, providing theoretical and empirical guarantees against planner collapse.

On real robot platforms (Franka FR3 and ARX X5 cube manipulation), PRISM achieves success rates of 35–45% in preliminary trials—substantially higher than direct-execution policies under deployment-induced goal distribution shifts. The lightweight design proves immediately transferable to hardware without algorithmic change and is essential for maintaining real-time responsiveness in embedded systems.

Implications and Limitations

By extracting the action prior from the same encoder and demonstration dataset as the dynamics predictor, PRISM eliminates architectural and computational bloat associated with separate policy or perceptual models, improving both efficiency and deployability. Its precision-weighted fusion lends itself to robust planner initialization, especially in regimes of limited compute where efficiency is critical for embedded and physical control systems. The unimodal prior suffices to outperform even large-scale multimodal policies, highlighting the value of principled initialization in sampling-based MPC.

PRISM is limited primarily by the expressivity and statistical coverage of its training data. It requires near-expert, task-specific demonstrations, restricting zero-shot generalization; performance on truly multimodal or long-horizon tasks may plateau given only unimodal prior heads. Future extensions toward mixture-of-Gaussians or diffusion-based priors could further improve planning performance in these regimes. Additionally, its efficacy in real-world hardware deployment remains to be quantified against matched baselines beyond proofs-of-concept.

Future Directions

The framework of precision-guided PoG fusion offers a wider principle for integrating learned priors from diverse sources into planners for both vision and language-conditioned tasks. Potential directions include adapting the fusion mechanism to reward- or value-based world models, developing richer multimodal priors, and exploring its applicability to hierarchical planners or alternative optimization algorithms (e.g., adaptive variance schemes as in CEM).

Conclusion

PRISM demonstrates that robust, state-conditioned action priors can be distilled directly from world model latent representations and fused into planner initialization via closed-form precision arithmetic. This approach eliminates redundancy, scales efficiently with compute, and preserves safety through graceful degradation. The method delivers strong empirical gains in both simulation and preliminary real-robot tests, underscoring its utility for practical continuous control in embodied AI. Future research will extend this fusion paradigm to richer prior models and broader planning architectures, advancing both theoretical understanding and practical deployment in AI planning (2606.07974).